From ca11b189b89badd8d53e880123046a34c314e685 Mon Sep 17 00:00:00 2001 From: Amog Kamsetty Date: Mon, 1 Mar 2021 20:14:35 -0800 Subject: [PATCH] [Tune] use epoch for ptl checkpoint dir name (#14392) * use epoch for dir name * use formatted string --- python/ray/tune/integration/pytorch_lightning.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/ray/tune/integration/pytorch_lightning.py b/python/ray/tune/integration/pytorch_lightning.py index 866b669ae..1f2fbfa53 100644 --- a/python/ray/tune/integration/pytorch_lightning.py +++ b/python/ray/tune/integration/pytorch_lightning.py @@ -217,7 +217,8 @@ class _TuneCheckpointCallback(TuneCallback): def _handle(self, trainer: Trainer, pl_module: LightningModule): if trainer.running_sanity_check: return - with tune.checkpoint_dir(step=trainer.global_step) as checkpoint_dir: + step = f"epoch={trainer.current_epoch}-step={trainer.global_step}" + with tune.checkpoint_dir(step=step) as checkpoint_dir: trainer.save_checkpoint( os.path.join(checkpoint_dir, self._filename))