Hi all, I have 2 list, A and B, the sample is a pair that consists of (a_i, b_j). Now the memory is not enough to fit N(A)xN(B) size data, so I want to generate one sample every time fetching it. My code is something like:

class MyDataset(Dataset):
    def __init__:
    ....
    def generate_one_sample(self):
        return (a_i,b_j)
    def __getitem__(self):
        return self.generate_one_sample()


This could work on CPU mode, BUT when I try to use GPU, this does not work, is there any solution for this case?
Thanks in advance