From 578a478924e29a2290e71d985e3685ab80ec3ebd Mon Sep 17 00:00:00 2001 From: devangpratap <115096812+devangpratap@users.noreply.github.com> Date: Mon, 15 Jun 2026 10:37:23 -0400 Subject: [PATCH] fix: use _get_autocast_context for fp32 action head (MPS compat) --- src/lerobot/policies/vla_jepa/modeling_vla_jepa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index a9c5290eb..0f97bb120 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -267,7 +267,7 @@ class VLAJEPAModel(nn.Module): ) -> Tensor: """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`.""" 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 horizon = self.config.chunk_size actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)