minor fix

datamol-io · Aug 11, 2023 · 4c08b56 · 4c08b56
1 parent bdf14d1
commit 4c08b56
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/graphium/data/datamodule.py b/graphium/data/datamodule.py
@@ -1191,7 +1191,7 @@ def setup(
         labels_size = {}
         labels_dtype = {}
         if stage == "fit" or stage is None:
-            if self.processed_graph_data_path is not None:
+            if self.dataloading_from == "disk":
                 processed_train_data_path = self._path_to_load_from_file("train")
                 assert self._data_ready_at_path(
                     processed_train_data_path

diff --git a/profiling/profile_predictor.py b/profiling/profile_predictor.py
@@ -20,7 +20,9 @@ def main():
 
     with fsspec.open(CONFIG_PATH, "r") as f:
         cfg = yaml.safe_load(f)
-    cfg["datamodule"]["args"]["processed_graph_data_path"] = "graphium/data/cache/profiling/predictor_data.cache"
+    cfg["datamodule"]["args"][
+        "processed_graph_data_path"
+    ] = "graphium/data/cache/profiling/predictor_data.cache"
     # cfg["datamodule"]["args"]["df_path"] = DATA_PATH
     cfg["trainer"]["trainer"]["max_epochs"] = 5
     cfg["trainer"]["trainer"]["min_epochs"] = 5