Dataloader fetch samples on the fly?

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

What kind of error do you get when using GPU?
Also, could you provide a small code snippet with the error?

Thanks. The memory left for me now is not enough to reproduce the error now, the process will be killed by OS due to out of memory, I will paste the error message when I get it.