admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-21 08:51:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fixing final upload to hub

Browse Source
This commit is contained in:
Maxime Ellerbach
2026-06-17 09:42:05 +00:00
parent b42d124007
commit 2e0deff3ab
2 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pretrained.py
+3 -1
View File
@@ -303,6 +303,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
self, self,
cfg: TrainPipelineConfig, cfg: TrainPipelineConfig,
peft_model=None, peft_model=None,
state_dict: dict[str, Tensor] | None = None,
): ):
api = HfApi() api = HfApi()
repo_id = api.create_repo( repo_id = api.create_repo(
@@ -320,7 +321,8 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
peft_model.save_pretrained(saved_path) peft_model.save_pretrained(saved_path)
self.config.save_pretrained(saved_path) self.config.save_pretrained(saved_path)
else: else:
self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors # Calls _save_pretrained and stores model tensors
self.save_pretrained(saved_path, state_dict=state_dict)
card = self.generate_model_card( card = self.generate_model_card(
cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg
src/lerobot/scripts/lerobot_train.py
+3 -1
View File
@@ -644,6 +644,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if eval_env: if eval_env:
close_envs(eval_env) close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process: if is_main_process:
logging.info("End of training") logging.info("End of training")
@@ -653,7 +655,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not cfg.is_reward_model_training and cfg.policy.use_peft: if not cfg.is_reward_model_training and cfg.policy.use_peft:
unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model)
else: else:
unwrapped_model.push_model_to_hub(cfg) unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict)
preprocessor.push_to_hub(active_cfg.repo_id) preprocessor.push_to_hub(active_cfg.repo_id)
postprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id)