bugfix with loading torch models

app.py

@@ -100,11 +100,20 @@ def load_data(batch_size):
     with open("saved_model/rf_class_train_tubulin.pickle", "rb") as inp:
         rf_model = pickle.load(inp)
 
-    vqgae_model = VQGAE.load_from_checkpoint("saved_model/vqgae.ckpt", task="decode", batch_size=batch_size)
-    vqgae_model = vqgae_model.to("cpu").eval()
-
-    ordering_model = OrderingNetwork.load_from_checkpoint("saved_model/ordering_network.ckpt", batch_size=batch_size)
-    ordering_model = ordering_model.to("cpu").eval()
+    vqgae_model = VQGAE.load_from_checkpoint(
+        "saved_model/vqgae.ckpt",
+        task="decode",
+        batch_size=batch_size,
+        map_location="cpu"
+    )
+    vqgae_model = vqgae_model.eval()
+
+    ordering_model = OrderingNetwork.load_from_checkpoint(
+        "saved_model/ordering_network.ckpt",
+        batch_size=batch_size,
+        map_location="cpu"
+    )
+    ordering_model = ordering_model.eval()
     return X, Y, rf_model, vqgae_model, ordering_model