fix rollout file name error

geoelements · Oct 24, 2024 · a97aa60 · a97aa60
1 parent 166470c
commit a97aa60
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/gns/train.py b/gns/train.py
@@ -200,8 +200,7 @@ def predict(device: str, cfg: DictConfig):
             if cfg.mode == "rollout":
                 example_rollout["metadata"] = metadata
                 example_rollout["loss"] = loss.mean()
-                filename = f"{cfg.output.filename}_ex{example_i}.pkl"
-                filename_render = f"{cfg.output.filename}_ex{example_i}"
+                filename_render = f"{cfg.output.filename}_ex{example_i}.pkl"
                 filename = os.path.join(cfg.output.path, filename_render)
                 with open(filename, "wb") as f:
                     pickle.dump(example_rollout, f)