Merge pull request #2 from weifengpy/fsdp2

enable saving checkpoint
weifengpy · May 1, 2024 · 7bbe522 · 7bbe522
2 parents b6fad93 + d6af9a2
commit 7bbe522
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/recipes/lora_finetune_distributed.py b/recipes/lora_finetune_distributed.py
@@ -589,7 +589,7 @@ def train(self) -> None:
                         )
 
                 self.epochs_run += 1
-                # self.save_checkpoint(epoch=curr_epoch)
+                self.save_checkpoint(epoch=curr_epoch)
 
     def cleanup(self) -> None:
         if self._is_rank_zero: