ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-31 10:51:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

adressing dtype zeros issue

Browse Source
This commit is contained in:
Maximellerbach
2026-05-21 15:40:17 +02:00
parent f75a2ee2f5
commit c01a00a972
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/lerobot/policies/vla_jepa/modeling_vla_jepa.py
View File

@@ -530,9 +530,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
examples = self._prepare_model_inputs(batch) examples = self._prepare_model_inputs(batch)
native_output = self.model.forward(examples) native_output = self.model.forward(examples)
total_loss = native_output.get("action_loss", torch.tensor(0.0)) + native_output.get( ref = next(iter(native_output.values()))
"wm_loss", torch.tensor(0.0) zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
) total_loss = native_output.get("action_loss", zero) + native_output.get("wm_loss", zero)
logs = {k: v.detach().item() for k, v in native_output.items()} logs = {k: v.detach().item() for k, v in native_output.items()}
logs["loss"] = total_loss.detach().item() logs["loss"] = total_loss.detach().item()
return total_loss, logs return total_loss, logs
@@ -577,6 +577,10 @@ class VLAJEPAPolicy(PreTrainedPolicy):
@classmethod @classmethod
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
"""
Custom loading to enable opt reinit of action head
when loading pretrained weights with mismatched action head shapes.
"""
if not model.config.reinit_action_head: if not model.config.reinit_action_head:
return super()._load_as_safetensor(model, model_file, map_location, strict) return super()._load_as_safetensor(model, model_file, map_location, strict)