diff --git a/recipes/lora_finetune_single_device.py b/recipes/lora_finetune_single_device.py index b0ffa49244..02b28da139 100644 --- 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: """