在训练中的过程中,会使用以下代码清空optimizer的梯度:
optimizer.zero_grad()
虽然在test过程中梯度不需要反向传播,但是网络还是会保存optimizer的梯度,如果optimizer保存在cuda中,显存就会逐渐增大。
解决方法:
with torch.no_grad():
for batch_idx, (img, label) in enumerate(testloader):
img, label = Variable(img).cuda(), Variable(label).cuda()
predict = model(img)