fix(autocast): route inference autocasts through safe helper

Apply get_safe_autocast_context to the control_utils and sync inference
paths for uniformity with lerobot_eval. AMP is now enabled on any
AMP-capable device (cuda, xpu, cpu) when use_amp is set, and stays a
no-op on mps.

docs/source/hardware_guide.mdx

@@ -82,18 +82,18 @@ VRAM is the first filter. Within a tier, pick by budget and availability — the
 
 ### Hugging Face Jobs
 
-[Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) lets you run training on managed HF infrastructure, billed by the second, without owning a GPU. `lerobot-train` submits and streams the job for you — just add `--job.target=<flavor>` to a normal training command:
+[Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) lets you run training on managed HF infrastructure, billed by the second. The repo publishes a ready-to-use image: **`huggingface/lerobot-gpu:latest`**, rebuilt **every night at 02:00 UTC from `main`** ([`docker_publish.yml`](https://github.com/huggingface/lerobot/blob/main/.github/workflows/docker_publish.yml)) — so it tracks the current state of the repo, not a tagged release.
 
 ```bash
-lerobot-train \
-  --policy.type=act --dataset.repo_id=<USER>/<DATASET> \
-  --policy.repo_id=<USER>/act_<task> \
-  --job.target=a10g-large
+hf jobs run --flavor a10g-large huggingface/lerobot-gpu:latest \
+  bash -c "nvidia-smi && lerobot-train \
+    --policy.type=act --dataset.repo_id=<USER>/<DATASET> \
+    --policy.repo_id=<USER>/act_<task> --batch_size=8 --steps=50000"
 ```
 
 Notes:
 
-- Run `hf auth login` once before submitting, the job runs under your token.
-- `--job.target` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). List the current catalogue with pricing via `hf jobs hardware`, or see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs).
-- The job defaults to a `2d` (48h) timeout. Override it with `--job.timeout=4h` (or any other valid duration string) to shorten or extend the timeout. The job automatically stops when the command completes.
-- For the full walkthrough — dataset upload, checkpoint streaming, resuming a run on a job — see the [imitation-learning training guide](./il_robots#train-using-hugging-face-jobs).
+- The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched.
+- The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training.
+- `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs).
+- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=<flavor>` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots).

docs/source/il_robots.mdx

@@ -532,7 +532,84 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola
 
 Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs).
 
-`lerobot-train` runs locally by default. To run on a HuggingFace GPU, pass `--job.target` with a hardware flavor name:
+> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs.
+
+To run the training manually use this command:
+
+<hfoptions id="train_with_hf_jobs">
+<hfoption id="Command">
+```bash
+hf jobs run \
+  --flavor a10g-small \
+  --timeout 4h \
+  --secrets HF_TOKEN \
+  huggingface/lerobot-gpu:latest \
+  -- \
+  python -m lerobot.scripts.lerobot_train \
+    --dataset.repo_id=username/dataset \
+    --policy.type=act \
+    --steps=5000 \
+    --batch_size=16 \
+    --policy.device=cuda \
+    --policy.repo_id=username/your_policy \
+    --log_freq=100
+```
+</hfoption>
+<hfoption id="API example">
+
+<!-- prettier-ignore-start -->
+```python
+from huggingface_hub import run_job, get_token
+
+run_name = "act_so101_hf_jobs"
+dataset_id = "username/dataset"
+user_hub_id = "username"
+
+command_args = [
+    "python", "-m", "lerobot.scripts.lerobot_train",
+    "--dataset.repo_id", dataset_id,
+    "--policy.type", "act",
+    "--steps", "5000",
+    "--batch_size", "16",
+    "--num_workers", "4",
+    "--policy.device", "cuda",
+    "--log_freq", "100",
+    "--save_freq", "1000",
+    "--save_checkpoint", "true",
+    "--wandb.enable", "false",
+    "--policy.repo_id", f"{user_hub_id}/{run_name}"
+]
+
+print(f"Submitting job '{run_name}' to Hugging Face Infrastructure...")
+
+job_info = run_job(
+    image="huggingface/lerobot-gpu:latest",
+    command=command_args,
+    flavor="a10g-small",
+    timeout="4h",
+    secrets={"HF_TOKEN": get_token()}
+)
+
+print("\n🚀 Job successfully launched!")
+print(f"🔹 Job ID: {job_info.id}")
+print(f"🔗 Live UI Dashboard & Logs: {job_info.url}")
+```
+<!-- prettier-ignore-end -->
+
+</hfoption>
+</hfoptions>
+
+You can modify the `--flavor` to use different hardware, for example: `t4-small`, `a100-large`, `h200`. Use `hf jobs hardware` to see the full list with pricing.
+Depending on the model you want to train and the hardware you selected you can also modify the `--batch_size` and `--number_of_workers`.
+For longer training sessions increase the timeout.
+
+Once the training is started you can go to [Jobs](https://huggingface.co/settings/jobs) and see if your jobs is running as well as all the outputs. Sometimes it takes a few minutes to schedule your job so be patient.
+
+After training the model will be pushed to hub and you can use it as any other model with LeRobot.
+
+#### Train on HF Jobs via `--job.target` (integrated CLI)
+
+`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name:
 
 ```bash
 lerobot-train \

src/lerobot/common/control_utils.py

@@ -18,7 +18,6 @@ from __future__ import annotations
 # Utilities
 ########################################################################################
 import time
-from contextlib import nullcontext
 from copy import copy
 from typing import TYPE_CHECKING, Any
 
@@ -26,6 +25,7 @@ import numpy as np
 import torch
 
 from lerobot.policies import PreTrainedPolicy, prepare_observation_for_inference
+from lerobot.utils.device_utils import get_safe_autocast_context
 from lerobot.utils.import_utils import _deepdiff_available, require_package
 
 if TYPE_CHECKING or _deepdiff_available:
@@ -76,7 +76,7 @@ def predict_action(
     observation = copy(observation)
     with (
         torch.inference_mode(),
-        torch.autocast(device_type=device.type) if device.type == "cuda" and use_amp else nullcontext(),
+        get_safe_autocast_context(device, enabled=use_amp),
     ):
         # Convert to pytorch format: channel first and float32 in [0,1] with batch dimension
         observation = prepare_observation_for_inference(observation, device, task, robot_type)

src/lerobot/datasets/compute_stats.py

@@ -519,13 +519,6 @@ def compute_episode_stats(
         if features[key]["dtype"] in {"string", "language"}:
             continue
 
-        # Features with zero-width shapes are skipped (no data to compute stats on)
-        if any(d == 0 for d in features[key].get("shape", ())):
-            logging.debug(
-                f"Skipping statistics computation for feature '{key}' with a zero-width shape {features[key]['shape']}."
-            )
-            continue
-
         if features[key]["dtype"] in ["image", "video"]:
             ep_ft_array = sample_images(data)
             axes_to_reduce = (0, 2, 3)

src/lerobot/datasets/feature_utils.py

@@ -67,9 +67,9 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
         elif ft["shape"] == (1,):
             hf_features[key] = datasets.Value(dtype=ft["dtype"])
         elif len(ft["shape"]) == 1:
-            # pyarrow rejects fixed-size lists of length 0, so use a variable length list instead
-            length = ft["shape"][0] if ft["shape"][0] > 0 else -1
-            hf_features[key] = datasets.Sequence(length=length, feature=datasets.Value(dtype=ft["dtype"]))
+            hf_features[key] = datasets.Sequence(
+                length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
+            )
         elif len(ft["shape"]) == 2:
             hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
         elif len(ft["shape"]) == 3:

src/lerobot/policies/groot/modeling_groot.py

@@ -43,6 +43,7 @@ from torch import Tensor
 
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.utils.constants import ACTION, OBS_IMAGES
+from lerobot.utils.device_utils import get_safe_autocast_context
 from lerobot.utils.import_utils import require_package
 
 from ..pretrained import PreTrainedPolicy
@@ -243,7 +244,7 @@ class GrootPolicy(PreTrainedPolicy):
 
         # Run GR00T forward under bf16 autocast when enabled to reduce activation memory
         # Rationale: Matches original GR00T finetuning (bf16 compute, fp32 params) and avoids fp32 upcasts.
-        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=self.config.use_bf16):
+        with get_safe_autocast_context(device, dtype=torch.bfloat16, enabled=self.config.use_bf16):
             outputs = self._groot_model.forward(groot_inputs)
 
         # Isaac-GR00T returns a BatchFeature; loss key is typically 'loss'
@@ -275,7 +276,7 @@ class GrootPolicy(PreTrainedPolicy):
         device = next(self.parameters()).device
 
         # Use bf16 autocast for inference to keep memory low and match backbone dtype
-        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=self.config.use_bf16):
+        with get_safe_autocast_context(device, dtype=torch.bfloat16, enabled=self.config.use_bf16):
             outputs = self._groot_model.get_action(groot_inputs)
 
         actions = outputs.get("action_pred")

src/lerobot/policies/molmoact2/modeling_molmoact2.py

@@ -31,7 +31,6 @@ import logging
 import os
 import types
 from collections import deque
-from contextlib import nullcontext
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -43,6 +42,7 @@ from torch.distributions import Beta
 
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
+from lerobot.utils.device_utils import get_safe_autocast_context
 from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
 
 from ..rtc.modeling_rtc import RTCProcessor
@@ -1644,10 +1644,8 @@ class MolmoAct2Policy(PreTrainedPolicy):
                 device=device,
             )
         action_dim = self._output_action_dim(batch)
-        autocast_context = (
-            torch.autocast(device_type=device.type, dtype=model_dtype)
-            if device.type in {"cuda", "cpu"} and model_dtype in {torch.bfloat16, torch.float16}
-            else nullcontext()
+        autocast_context = get_safe_autocast_context(
+            device, dtype=model_dtype, enabled=model_dtype in {torch.bfloat16, torch.float16}
         )
         with autocast_context:
             if inference_action_mode == "discrete":

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -26,6 +26,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 get_safe_autocast_context
 from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
@@ -183,7 +184,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_safe_autocast_context(device_type, dtype=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 +251,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_safe_autocast_context(device_type, dtype=torch.float32):
             r = self.config.repeated_diffusion_steps
             horizon = self.config.chunk_size
             actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)

src/lerobot/rollout/inference/sync.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import logging
-from contextlib import nullcontext
 from copy import copy
 
 import torch
@@ -25,6 +24,7 @@ import torch
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
 from lerobot.processor import PolicyProcessorPipeline
+from lerobot.utils.device_utils import get_safe_autocast_context
 
 from .base import InferenceEngine
 
@@ -102,11 +102,7 @@ class SyncInferenceEngine(InferenceEngine):
         # ``obs_frame`` fresh per tick via ``build_dataset_frame``, so the
         # tensor/array values are not shared with any other reader.
         observation = copy(obs_frame)
-        autocast_ctx = (
-            torch.autocast(device_type=self._device.type)
-            if self._device.type == "cuda" and self._policy.config.use_amp
-            else nullcontext()
-        )
+        autocast_ctx = get_safe_autocast_context(self._device, enabled=self._policy.config.use_amp)
         with torch.inference_mode(), autocast_ctx:
             observation = prepare_observation_for_inference(
                 observation, self._device, self._task, self._robot_type

src/lerobot/scripts/lerobot_eval.py

@@ -56,7 +56,6 @@ import threading
 import time
 from collections import defaultdict
 from collections.abc import Callable
-from contextlib import nullcontext
 from copy import deepcopy
 from dataclasses import asdict
 from functools import partial
@@ -86,7 +85,7 @@ from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_proces
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
-from lerobot.utils.device_utils import get_safe_torch_device
+from lerobot.utils.device_utils import get_safe_autocast_context, get_safe_torch_device
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.io_utils import write_video
 from lerobot.utils.random_utils import set_seed
@@ -698,7 +697,7 @@ def eval_main(cfg: EvalPipelineConfig):
     max_episodes_rendered = 0 if cfg.eval.recording else 10
     videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos"
 
-    with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
+    with torch.no_grad(), get_safe_autocast_context(device, enabled=cfg.policy.use_amp):
         info = eval_policy_all(
             envs=envs,
             policy=policy,

src/lerobot/utils/__init__.py

@@ -33,7 +33,12 @@ from .constants import (
     REWARD,
 )
 from .decorators import check_if_already_connected, check_if_not_connected
-from .device_utils import auto_select_torch_device, get_safe_torch_device, is_torch_device_available
+from .device_utils import (
+    auto_select_torch_device,
+    get_safe_autocast_context,
+    get_safe_torch_device,
+    is_torch_device_available,
+)
 from .errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from .import_utils import is_package_available, require_package
 
@@ -51,6 +56,7 @@ __all__ = [
     "REWARD",
     # Device utilities
     "auto_select_torch_device",
+    "get_safe_autocast_context",
     "get_safe_torch_device",
     "is_torch_device_available",
     # Import guards

src/lerobot/utils/device_utils.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import logging
+from contextlib import AbstractContextManager, nullcontext
 
 import torch
 
@@ -107,3 +108,25 @@ def is_amp_available(device: str):
         return False
     else:
         raise ValueError(f"Unknown device '{device}.")
+
+
+def get_safe_autocast_context(
+    device: str | torch.device,
+    *,
+    dtype: torch.dtype | None = None,
+    enabled: bool = True,
+) -> AbstractContextManager:
+    """Return a ``torch.autocast`` context, or a no-op when AMP is unsupported on ``device``.
+
+    Autocast is only entered on devices that support AMP (cuda, xpu, cpu); on mps and any
+    unknown device this falls back to ``nullcontext()`` so callers can request autocast
+    unconditionally without breaking on unsupported backends.
+    """
+    device_type = device.type if isinstance(device, torch.device) else str(device).split(":", 1)[0]
+    try:
+        amp_ok = is_amp_available(device_type)
+    except ValueError:
+        amp_ok = False
+    if not enabled or not amp_ok:
+        return nullcontext()
+    return torch.autocast(device_type=device_type, dtype=dtype)

tests/datasets/test_compute_stats.py

@@ -13,7 +13,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import logging
 from unittest.mock import patch
 
 import numpy as np
@@ -688,28 +687,6 @@ def test_compute_episode_stats_string_features_skipped():
     assert "q01" in stats["action"]
 
 
-def test_compute_episode_stats_zero_width_features_skipped(caplog):
-    """Test that features with a zero-width dim (e.g. shape=(0,)) are skipped with a debug log."""
-    episode_data = {
-        "empty": np.zeros((100, 0), dtype=np.float32),  # Zero-width feature
-        "action": np.random.normal(0, 1, (100, 5)),
-    }
-    features = {
-        "empty": {"dtype": "float32", "shape": (0,)},
-        "action": {"dtype": "float32", "shape": (5,)},
-    }
-
-    with caplog.at_level(logging.DEBUG):
-        stats = compute_episode_stats(episode_data, features)
-
-    # Zero-width features should be skipped with a debug log, others computed as usual
-    assert "empty" not in stats
-    assert "empty" in caplog.text
-    assert "action" in stats
-    assert "q01" in stats["action"]
-    assert stats["action"]["mean"].shape == (5,)
-
-
 def test_aggregate_feature_stats_with_quantiles():
     """Test aggregating feature stats that include quantiles."""
     stats_ft_list = [

tests/datasets/test_datasets.py

@@ -1804,11 +1804,3 @@ def test_episode_filter_unknown_key_raises(tmp_path, lerobot_dataset_factory):
             root=dataset.root,
             episode_filter=lambda ep: ep["not_a_real_field"] > 0,
         )
-
-
-def test_get_hf_features_zero_width_feature_does_not_raise_on_from_dict():
-    import datasets
-
-    features = {"empty": {"dtype": "float32", "shape": (0,), "names": ["empty"]}}
-    hf_features = get_hf_features_from_features(features)
-    datasets.Dataset.from_dict({"empty": [[], []]}, features=hf_features)

tests/utils/test_device_utils.py

@@ -0,0 +1,40 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+
+from lerobot.utils.device_utils import get_safe_autocast_context
+
+
+@pytest.mark.parametrize(
+    ("device", "enabled", "expect_autocast"),
+    [
+        ("cpu", True, True),  # AMP-capable device -> real autocast
+        (torch.device("cpu"), True, True),  # accepts torch.device
+        ("cpu", False, False),  # explicitly disabled -> no-op
+        ("mps", True, False),  # AMP unsupported on mps -> no-op
+        ("privateuseone", True, False),  # unknown device -> safe no-op
+    ],
+)
+def test_get_safe_autocast_context(device, enabled, expect_autocast):
+    ctx = get_safe_autocast_context(device, dtype=torch.bfloat16, enabled=enabled)
+    if expect_autocast:
+        assert isinstance(ctx, torch.autocast)
+        with ctx:
+            assert torch.is_autocast_enabled("cpu")
+    else:
+        assert isinstance(ctx, nullcontext)