add flag to store all checkpoints

src/gflownet/config.py

@@ -60,6 +60,8 @@ class Config:
         The number of training steps after which to validate the model
     checkpoint_every : Optional[int]
         The number of training steps after which to checkpoint the model
+    store_all_checkpoints : bool
+        Whether to store all checkpoints or only the last one
     print_every : int
         The number of training steps after which to print the training loss
     start_at_step : int
@@ -85,6 +87,7 @@ class Config:
     seed: int = 0
     validate_every: int = 1000
     checkpoint_every: Optional[int] = None
+    store_all_checkpoints: bool = False
     print_every: int = 100
     start_at_step: int = 0
     num_final_gen_steps: Optional[int] = None

src/gflownet/trainer.py

@@ -418,7 +418,8 @@ def _save_state(self, it):
                 state,
                 fd,
             )
-        shutil.copy(fn, pathlib.Path(self.cfg.log_dir) / f"model_state_{it}.pt")
+        if self.cfg.store_all_checkpoints:
+            shutil.copy(fn, pathlib.Path(self.cfg.log_dir) / f"model_state_{it}.pt")
 
     def log(self, info, index, key):
         if not hasattr(self, "_summary_writer"):