方法1:batch-size设置多小
方法2:
with torch.no_grad():
net = Net()
out = net(imgs)
积累的梯度应该是会一直放在显存里的...用了这一行就会停止自动反向计算梯度
方法3:
设置cpu来加载模型:
model_path = 'path/to/model.pt'
model = UNet(n_channels = 1, n_classes = 1)
state_dict = torch.load(model_path,map_location='cpu')
model.load_state_dict(state_dict)
model.to(device)