Update yolo.py

yolo.py

@@ -67,8 +67,9 @@ class YOLO(object):
 
         # 加快模型训练的效率
         print('Loading weights into state dict...')
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         state_dict = torch.load(self.model_path)
-        self.net.load_state_dict(state_dict)
+        self.net.load_state_dict(state_dict, map_location=device)
         
         if self.cuda:
             os.environ["CUDA_VISIBLE_DEVICES"] = '0'