remove loading model on CPU by default

d-f · Mar 18, 2024 · a3bc1ba · a3bc1ba
1 parent 71e5d10
commit a3bc1ba
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train_torchvision.py b/train_torchvision.py
@@ -104,7 +104,7 @@ def load_model(weight_path: Path, model: models.efficientnet.EfficientNet) -> mo
     """
     loads all parameters of a model
     """
-    checkpoint = torch.load(weight_path, map_location='cpu')
+    checkpoint = torch.load(weight_path)
     model.load_state_dict(checkpoint['state_dict'])
     return model