admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-27 19:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: use _get_autocast_context for fp32 action head (MPS compat)

Browse Source
This commit is contained in:
devangpratap
2026-06-15 10:37:23 -04:00
committed by Maxime Ellerbach
parent d648e308d2
commit 578a478924
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/vla_jepa/modeling_vla_jepa.py
+1 -1
View File
@@ -267,7 +267,7 @@ class VLAJEPAModel(nn.Module):
) -> Tensor: ) -> Tensor:
"""Flow-matching action-head loss, repeated over `repeated_diffusion_steps`.""" """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
device_type = next(self.parameters()).device.type device_type = next(self.parameters()).device.type
with torch.autocast(device_type=device_type, dtype=torch.float32): with _get_autocast_context(device_type, torch.float32):
r = self.config.repeated_diffusion_steps r = self.config.repeated_diffusion_steps
horizon = self.config.chunk_size horizon = self.config.chunk_size
actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1) actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)