load to cpu

src/rstor/analyzis/interactive/model_selection.py

@@ -26,7 +26,7 @@ def model_selector(models_dict: dict, global_params={}, model_name="vanilla"):
 def get_model_from_exp(exp: int, model_storage: Path = MODELS_PATH, device=DEVICE) -> Tuple[torch.nn.Module, dict]:
     config = get_experiment_config(exp)
     model, _, _ = get_training_content(config, training_mode=False)
-    model_path = torch.load(model_storage/f"{exp:04d}"/"best_model.pt")
+    model_path = torch.load(model_storage/f"{exp:04d}"/"best_model.pt", map_location=device)
     assert model_path is not None, f"Model {exp} not found"
     model.load_state_dict(model_path)
     model = model.to(device)