fix(vla_jepa): use device-safe autocast instead of hardcoded bfloat16 (#4158)

* fix(vla_jepa): use device-safe autocast instead of hardcoded bfloat16

VLA-JEPA hardcodes torch.autocast with dtype=torch.bfloat16, which
crashes on MPS (no AMP support) and silently misbehaves on pre-Ampere
CUDA GPUs (no bf16). Add a _get_autocast_context() helper that reuses
the existing is_amp_available() utility to pick a safe strategy per
device, matching the pattern used by pi05 and molmoact2.

Fixes #3744

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

---------

Co-authored-by: devangpratap <115096812+devangpratap@users.noreply.github.com>

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import logging
 from collections import deque
+from contextlib import nullcontext
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -26,6 +27,7 @@ from torch import Tensor, nn
 from lerobot.policies.pretrained import PreTrainedPolicy, T
 from lerobot.policies.utils import populate_queues
 from lerobot.utils.constants import ACTION, OBS_STATE
+from lerobot.utils.device_utils import is_amp_available
 from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
@@ -39,6 +41,21 @@ from .configuration_vla_jepa import VLAJEPAConfig
 from .qwen_interface import Qwen3VLInterface
 from .world_model import ActionConditionedVideoPredictor
 
+
+def _get_autocast_context(device_type: str, dtype: torch.dtype = torch.bfloat16):
+    """Return an autocast context appropriate for the device.
+
+    MPS does not support ``torch.autocast`` at all.  On CUDA devices
+    without bfloat16 support (compute capability < 8.0) we fall back to
+    float16.
+    """
+    if not is_amp_available(device_type):
+        return nullcontext()
+    if device_type == "cuda" and dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
+        dtype = torch.float16
+    return torch.autocast(device_type=device_type, dtype=dtype)
+
+
 # ============================================================================
 # Native VLA-JEPA Model - follows original starVLA VLA_JEPA.py implementation
 # ============================================================================
@@ -183,7 +200,7 @@ class VLAJEPAModel(nn.Module):
             action_idx = action_mask.nonzero(as_tuple=True)
 
         device_type = next(self.parameters()).device.type
-        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+        with _get_autocast_context(device_type, torch.bfloat16):
             last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
             b, _, h = last_hidden.shape
             embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h)
@@ -250,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)