NumpyDataset for pytorch (like tensordataset)

numpy_dataset.py

import torch.utils.data


class NumpyDataset(torch.utils.data.Dataset):
    """Dataset wrapping arrays.

    Each sample will be retrieved by indexing array along the first dimension.

    Arguments:
        *arrays (numpy.array): arrays that have the same size of the first dimension.
    """

    def __init__(self, *arrays):
        assert all(arrays[0].shape[0] == array.shape[0] for array in arrays)
        self.arrays = arrays

    def __getitem__(self, index):
        return tuple(array[index].compute() for array in self.arrays)

    def __len__(self):
        return self.arrays[0].shape[0]