Skip to content

Commit e1b172c

Browse files
Profile Trainer.save_checkpoint (#20405)
1 parent f5b513d commit e1b172c

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

src/lightning/pytorch/trainer/trainer.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1361,9 +1361,10 @@ def save_checkpoint(
13611361
"Saving a checkpoint is only possible if a model is attached to the Trainer. Did you call"
13621362
" `Trainer.save_checkpoint()` before calling `Trainer.{fit,validate,test,predict}`?"
13631363
)
1364-
checkpoint = self._checkpoint_connector.dump_checkpoint(weights_only)
1365-
self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options)
1366-
self.strategy.barrier("Trainer.save_checkpoint")
1364+
with self.profiler.profile("save_checkpoint"):
1365+
checkpoint = self._checkpoint_connector.dump_checkpoint(weights_only)
1366+
self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options)
1367+
self.strategy.barrier("Trainer.save_checkpoint")
13671368

13681369
"""
13691370
State properties

0 commit comments

Comments
 (0)