Thanks for your advice!
Here’s the dataset that I’m using:
# fake data to show the problem
x = torch.cuda.FloatTensor(10000,3)
x = torch.randn(10000,3)
y = torch.cuda.FloatTensor(10000,3)
y = torch.randn(10000,3)

train = data_utils.TensorDataset(x, y)
train_loader = data_utils.DataLoader(train, batch_size=batchSize, shuffle=False)

I try to set the batchSize to 1, but the GPU memory is still keeping increasing during training.