From theoretical perspective, pytorch should not squeeze all training data into memory. There could be a bug somewhere in the pipeline, thanks for posting this, and will track this post. If I were you, I would do this: instead of using data loader, I can directly load image and transfer it to a tensor and use the tensor in the training loop. After doing this, if the memory issue goes away, that means the bug exists in the dataset/data loader side. If the memory still climbs up, I will go to check if the image data gets retained somewhere…