Fix cuda bug

app.py

@@ -52,7 +52,7 @@ def convert_models_to_fp32(model):
 
 
 def load(model_path, device, model, image_resolution):
-    state_dict = torch.load(model_path, map_location="cpu")
+    state_dict = torch.load(model_path, map_location=device)
     state_dict = state_dict["state_dict"]
 
     model_config_file = f"{model.replace('/', '-')}.json"
@@ -383,7 +383,7 @@ def molecules_from_image():
             predefined_features = False
         else:
             mol_index = pd.read_csv("cellpainting-unique-molecule.csv")
-            mol_features_torch = torch.load("all_molecule_cellpainting_features.pkl")
+            mol_features_torch = torch.load("all_molecule_cellpainting_features.pkl", map_location=device)
             mol_features = mol_features_torch["mol_features"]
             mol_ids = mol_features_torch["mol_ids"]
             print(len(mol_ids))
@@ -453,7 +453,7 @@ def images_from_molecule():
             st.write("")
 
 
-        img_features_torch = torch.load(image_features)
+        img_features_torch = torch.load(image_features, map_location=device)
         img_features = img_features_torch["img_features"]
         img_ids = img_features_torch["img_ids"]