Update train.py

wyf763240785 · Jun 4, 2020 · fb8d176 · fb8d176
1 parent 83a0995
commit fb8d176
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -77,8 +77,9 @@ def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epo
     Cuda = True
 
     print('Loading weights into state dict...')
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model_dict = model.state_dict()
-    pretrained_dict = torch.load("model_data/yolo_weights.pth")
+    pretrained_dict = torch.load("model_data/yolo_weights.pth", map_location=device)
     pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
     model_dict.update(pretrained_dict)
     model.load_state_dict(model_dict)