Fixed try block in load_from_path()

kazewong · Sep 2, 2024 · 337c21e · 337c21e
1 parent 7309e44
commit 337c21e
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/src/jimgw/single_event/runManager.py b/src/jimgw/single_event/runManager.py
@@ -123,17 +123,13 @@ def save(self, path: str):
     def load_from_path(self, path: str) -> SingleEventRun:
         with open(path, "r") as f:
             data = yaml.safe_load(f)
-        if (
-            "jim_parameters" in data
-            and "local_sampler_arg" in data["jim_parameters"]
-            and "step_size" in data["jim_parameters"]["local_sampler_arg"]
-        ):
             try:
                 data["jim_parameters"]["local_sampler_arg"]["step_size"] = jnp.array(
                     data["jim_parameters"]["local_sampler_arg"]["step_size"]
                 )
-            except Exception as e:
-                print(f"Error in loading step_size: {e}")
+            except KeyError as e:
+                print(f"Key {e} not found.")
+
         return SingleEventRun(**data)
 
     ### Initialization functions ###