Upd

pytorch · Apr 10, 2024 · 0b4fbe1 · 0b4fbe1
1 parent 0ebd6f6
commit 0b4fbe1
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/recipes/lora_finetune_single_device.py b/recipes/lora_finetune_single_device.py
@@ -374,6 +374,11 @@ def save_checkpoint(self, epoch: int) -> None:
             epoch=epoch,
             intermediate_checkpoint=(epoch + 1 < self.total_epochs),
         )
+        log.info(
+            utils.memory_stats_log(
+                "Memory Stats during checkpoint save:", device=self._device
+            )
+        )
 
     def train(self) -> None:
         """