hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 02:21:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Tune] use epoch for ptl checkpoint dir name (#14392)

Browse source 
* use epoch for dir name

* use formatted string
This commit is contained in:
Amog Kamsetty 2021-03-01 20:14:35 -08:00 • committed by GitHub
parent dbaa28f81e
commit ca11b189b8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
python/ray/tune/integration/pytorch_lightning.py
View file

@ -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))