pytorch 训练/测试模型时错误：RuntimeError: CUDA error: out of memory

方法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)