Make dir_save_path default to None

mllam · Dec 10, 2024 · e61bdfe · e61bdfe
1 parent 5ed304a
commit e61bdfe
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/neural_lam/build_rectangular_graph.py b/neural_lam/build_rectangular_graph.py
@@ -253,7 +253,7 @@ def build_graph_from_archetype(
     datastore_boundary,
     graph_name,
     archetype,
-    dir_save_path,
+    dir_save_path=None,
     **kwargs,
 ):
     """

diff --git a/tests/test_training.py b/tests/test_training.py
@@ -66,17 +66,16 @@ def test_training(datastore_name, datastore_boundary_name):
 
     graph_dir_path = Path(datastore.root_path) / "graphs" / flat_graph_name
 
-    def _create_graph():
-        if not graph_dir_path.exists():
-            build_graph_from_archetype(
-                datastore=datastore,
-                datastore_boundary=datastore_boundary,
-                graph_name=flat_graph_name,
-                archetype="keisler",
-                mesh_node_distance=get_test_mesh_dist(
-                    datastore, datastore_boundary
-                ),
-            )
+    if not graph_dir_path.exists():
+        build_graph_from_archetype(
+            datastore=datastore,
+            datastore_boundary=datastore_boundary,
+            graph_name=flat_graph_name,
+            archetype="keisler",
+            mesh_node_distance=get_test_mesh_dist(
+                datastore, datastore_boundary
+            ),
+        )
 
     data_module = WeatherDataModule(
         datastore=datastore,