pytorch 清楚缓存_如何在PyTorch中清除Cuda内存

I am trying to get the output of a neural network which I have already trained. The input is an image of the size 300x300. I am using a batch size of 1, but I still get a CUDA error: out of memory error after I have successfully got the output for 25 images.

I searched for some solutions online and came across torch.cuda.empty_cache(). But this still doesn't seem to solve the problem.

This is the code I am using.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

train_x = torch.tensor(train_x, dtype=torch.float32).view(-1, 1, 300, 300)

train_x = train_x.to(device)

dataloader = torch.utils.data.DataLoader(train_x, batch_size=1, shuffle=False)

right = []

for i, left in enumerate(dataloader):

print(i)

temp = model(left).view(-1, 1, 300, 300)

right.append(temp.to('cpu'))