wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

backup wip

This commit is contained in:
Alexander Soare 2024-04-15 10:33:00 +01:00
parent 3df7f1d9b2
commit c362dc94e4
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
lerobot/common/logger.py
View File

@ -71,7 +71,12 @@ class Logger:
if self._save_model: if self._save_model:
self._model_dir.mkdir(parents=True, exist_ok=True) self._model_dir.mkdir(parents=True, exist_ok=True)
fp = self._model_dir / f"{str(identifier)}.pt" fp = self._model_dir / f"{str(identifier)}.pt"
policy.save_pretrained(fp) # TODO(alexander-soare): This conditional branching is temporary while we add PyTorchModelHubMixin
# to all policies.
if hasattr(policy, "save"):
policy.save(fp)
else:
policy.save_pretrained(fp)
if self._wandb and not self._disable_wandb_artifact: if self._wandb and not self._disable_wandb_artifact:
# note wandb artifact does not accept ":" in its name # note wandb artifact does not accept ":" in its name
artifact = self._wandb.Artifact( artifact = self._wandb.Artifact(