修改模型权重加载方式为'cpu'而不是'cuda'

utils.py

@@ -17,8 +17,8 @@ logger = logging
 
 def load_checkpoint(checkpoint_path, model, optimizer=None):
   assert os.path.isfile(checkpoint_path)
-  # checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-  checkpoint_dict = torch.load(checkpoint_path, map_location='cuda:0', weights_only=True)
+  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+  #checkpoint_dict = torch.load(checkpoint_path, map_location='cuda:0', weights_only=True)
 
   iteration = checkpoint_dict['iteration']
   learning_rate = checkpoint_dict['learning_rate']