# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/misc/torch.dataset.utils.ipynb.

# %% auto 0
__all__ = ['ListDataset', 'ZippedDataset', 'ZippedDataLoader']

# %% ../../../nbs/misc/torch.dataset.utils.ipynb 3
from torch.utils.data import Dataset

# %% ../../../nbs/misc/torch.dataset.utils.ipynb 4
class ListDataset(Dataset):
     def __init__(self, original_list):
        self.original_list = original_list
     
     def __len__(self):
        return len(self.original_list)

     def __getitem__(self, i):
        return self.original_list[i]

# %% ../../../nbs/misc/torch.dataset.utils.ipynb 5
class ZippedDataset:
    def __init__(self, *datasets):
        self.datasets = datasets
    
    def __getitem__(self, idx):
        return tuple(ds[idx] for ds in self.datasets)
    
    def __len__(self):
        return min(len(ds) for ds in self.datasets)

# %% ../../../nbs/misc/torch.dataset.utils.ipynb 6
class ZippedDataLoader:
    def __init__(self, *data_loaders):
        self.data_loaders = data_loaders

    def __iter__(self):
        self.it = iter(zip(*self.data_loaders))
        return self

    def __next__(self):
        return next(self.it)
