Hello, I am using GoogleColab to train a generative model and find that moving the data to GPU is the bottleneck of the training. Since Colab provides sufficient RAM, I did the following:

put all the training images in one NumPy array
create a tensor from it
move the tensor to GPU
create a TensorDataset from the tensor
feed TensorDataset to dataloader

As you can guess the cost of moving data to GPU is significantly decreased, while the inference step starts to cost almost the same amount of time, and what is even weird is that at the beginning of each epoch, the inference is almost the same speed, and within 3 batches it totally slows down…