forcing torch.load to CPU

app.py

@@ -9,7 +9,9 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 _old_load = torch.load
 
 def safe_torch_load(*args, **kwargs):
-    if 'map_location' not in kwargs:
+    if len(args) >= 2:
+        args[1] = device
+    else:
         kwargs['map_location'] = device
     return _old_load(*args, **kwargs)