Hi,
thanks for your answer. The reason I think it’s related to the DataLoader is that the memory used it directly correlated with the batch size. I tried doing del loss, del output, but didn’t see any changes. Here is my code if you wouldn’t mind taking a look  https://github.com/pclucas14/pixel-cnn-pp/blob/master/main.py 5
Thanks!