Save ckpt after every epoch.

bioscan-ml · Jul 13, 2024 · f371325 · f371325
1 parent fe91901
commit f371325
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/scripts/train_cl.py b/scripts/train_cl.py
@@ -144,10 +144,10 @@ def main_process(rank: int, world_size: int, args):
                 wandb.log(dict_for_wandb,
                           commit=True)
 
-            if args.save_ckpt:
-                last_ckpt_path = os.path.join(folder_path, f'last.pth')
-                torch.save(model.state_dict(), last_ckpt_path)
-                print(f'Last ckpt: {last_ckpt_path}')
+        if args.save_ckpt:
+            last_ckpt_path = os.path.join(folder_path, f'last.pth')
+            torch.save(model.state_dict(), last_ckpt_path)
+            print(f'Last ckpt: {last_ckpt_path}')
 
 
 @hydra.main(config_path="../bioscanclip/config", config_name="global_config", version_base="1.1")

diff --git a/scripts/train_cl_on_insect_dataset.py b/scripts/train_cl_on_insect_dataset.py
@@ -162,10 +162,10 @@ def main_process(rank: int, world_size: int, args):
                 wandb.log(dict_for_wandb,
                           commit=True)
 
-            if args.save_ckpt:
-                last_ckpt_path = os.path.join(folder_path, f'last_device_{rank}.pth')
-                torch.save(model.state_dict(), last_ckpt_path)
-                print(f'Last ckpt: {last_ckpt_path}')
+        if args.save_ckpt:
+            last_ckpt_path = os.path.join(folder_path, f'last_device_{rank}.pth')
+            torch.save(model.state_dict(), last_ckpt_path)
+            print(f'Last ckpt: {last_ckpt_path}')