解决Pytorch 训练与测试时爆显存(out of memory)的问题

try:
  output = model(input)
except RuntimeError as exception:
  if "out of memory" in str(exception):
    print("WARNING: out of memory")
    if hasattr(torch.cuda, 'empty_cache'):
      torch.cuda.empty_cache()
  else:
    raise exception

  with torch.no_grad():
    for ii,(inputs,filelist) in tqdm(enumerate(test_loader), desc='predict'):
      if opt.use_gpu:
        inputs = inputs.cuda()
        if len(inputs.shape) < 4:
          inputs = inputs.unsqueeze(1)
 
      else:
        if len(inputs.shape) < 4:
          inputs = torch.transpose(inputs, 1, 2)
          inputs = inputs.unsqueeze(1)