tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Checkpoint on final step of training even when it doesn't coincide with save_freq. (#284)

Browse Source
This commit is contained in:
Alexander Soare
2024-06-20 08:27:01 +01:00
committed by GitHub
parent 2abef3bef9
commit 9aa4cdb976
2 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
lerobot/scripts/train.py
View File

@@ -351,7 +351,10 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
logger.log_video(eval_info["video_paths"][0], step, mode="eval")
logging.info("Resume training")
if cfg.training.save_checkpoint and step % cfg.training.save_freq == 0:
if cfg.training.save_checkpoint and (
step % cfg.training.save_freq == 0
or step == cfg.training.offline_steps + cfg.training.online_steps
):
logging.info(f"Checkpoint policy after step {step}")
# Note: Save with step as the identifier, and format it to have at least 6 digits but more if
# needed (choose 6 as a minimum for consistency without being overkill).