Update TractionModel.py

app/TractionModel.py

@@ -53,7 +53,7 @@ def create_model():
     return model
 
 
-def load_weights(model, path='model.pt'):
-    checkpoint = torch.load(path, map_location=torch.device('cpu'))
+def load_weights(model, path='model.pt', device_='cpu'):
+    checkpoint = torch.load(path, map_location=torch.device(device_))
     model.load_state_dict(checkpoint)
     return model