From d17e4e5d4f7936ec5a2ce24282745c9b0447db0d Mon Sep 17 00:00:00 2001 From: Eric Cousineau Date: Thu, 7 Nov 2024 18:16:18 -0500 Subject: [PATCH] Profile Trainer.save_checkpoint --- src/lightning/pytorch/trainer/trainer.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index 406f686efe732..23db90fd45b38 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/src/lightning/pytorch/trainer/trainer.py @@ -1361,9 +1361,10 @@ def save_checkpoint( "Saving a checkpoint is only possible if a model is attached to the Trainer. Did you call" " `Trainer.save_checkpoint()` before calling `Trainer.{fit,validate,test,predict}`?" ) - checkpoint = self._checkpoint_connector.dump_checkpoint(weights_only) - self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options) - self.strategy.barrier("Trainer.save_checkpoint") + with self.profiler.profile("save_checkpoint"): + checkpoint = self._checkpoint_connector.dump_checkpoint(weights_only) + self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options) + self.strategy.barrier("Trainer.save_checkpoint") """ State properties