Update TractionModel.py

app/TractionModel.py

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