yichao_Liu:

Thing is this dataset is a customed dataset. I wrote a dataloader for it (not Pytorch DataLoader), but I think maybe it’s the problem with that I load all the dataset in init function.


It doesn’t really matter that you use your on dataloader. You just have to make sure that you only move one batch at a time on to the GPU. Like I did here:



 RaLo4:

    for i, data in enumerate(train_loader, 0):
        """training loop for one batch"""
        #get images and labels
        inputs, labels = data
        #move data to gpu
        inputs, labels = inputs.to(device), labels.to(device)



As long as inputs and labels is always just one batch it should which kind of dataloader they come from.
Also! Check if you are maybe putting different tensors onto the GPU and keeping them there.
If you for example do total_loss += loss  instead of  total_loss += loss.item() for logging your loss, your GPU memory will fill up.