fix(train): pick the local training device instead of assuming cuda

frontend/src/components/training/config/AdvancedCard.tsx

@@ -56,25 +56,31 @@ const AdvancedCard: React.FC<ConfigComponentProps> = ({ config, updateConfig })
           <section className="space-y-4">
             <SectionHeading>Policy</SectionHeading>
             <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
-              <div>
-                <Label htmlFor="policy_device" className="text-slate-300">
-                  Device
-                </Label>
-                <Select
-                  value={config.policy_device || 'cuda'}
-                  onValueChange={(value) => updateConfig('policy_device', value)}
-                >
-                  <SelectTrigger id="policy_device" className="bg-slate-900 border-slate-600 text-white rounded-lg">
-                    <SelectValue />
-                  </SelectTrigger>
-                  <SelectContent className="bg-slate-800 border-slate-600 text-white">
-                    <SelectItem value="cuda">CUDA (GPU)</SelectItem>
-                    <SelectItem value="cpu">CPU</SelectItem>
-                    <SelectItem value="mps">MPS (Apple Silicon)</SelectItem>
-                  </SelectContent>
-                </Select>
-              </div>
-              <div className="flex items-center space-x-3 pt-6">
+              {/* An HF Jobs pod picks its own device, so this only applies locally. */}
+              {config.target.runner === 'local' && (
+                <div>
+                  <Label htmlFor="policy_device" className="text-slate-300">
+                    Device
+                  </Label>
+                  <Select
+                    value={config.policy_device ?? 'auto'}
+                    onValueChange={(value) =>
+                      updateConfig('policy_device', value === 'auto' ? undefined : value)
+                    }
+                  >
+                    <SelectTrigger id="policy_device" className="bg-slate-900 border-slate-600 text-white rounded-lg">
+                      <SelectValue />
+                    </SelectTrigger>
+                    <SelectContent className="bg-slate-800 border-slate-600 text-white">
+                      <SelectItem value="auto">Auto</SelectItem>
+                      <SelectItem value="cpu">CPU</SelectItem>
+                    </SelectContent>
+                  </Select>
+                </div>
+              )}
+              <div
+                className={`flex items-center space-x-3 ${config.target.runner === 'local' ? 'pt-6' : ''}`}
+              >
                 <Switch
                   id="policy_use_amp"
                   checked={config.policy_use_amp}

frontend/src/pages/Training.tsx

@@ -134,7 +134,6 @@ const ConfigurationMode: React.FC = () => {
     wandb_enable: false,
     wandb_mode: "online",
     wandb_disable_artifact: false,
-    policy_device: "cuda",
     policy_use_amp: false,
     optimizer_type: "adam",
     use_policy_training_preset: true,

lelab/rollout.py

@@ -37,6 +37,8 @@
 
 from pydantic import BaseModel, field_validator
 
+from lerobot.utils.device_utils import auto_select_torch_device
+
 from .camera_config import lerobot_camera_settings, validate_camera_rotations
 from .utils.config import setup_follower_calibration_file
 
@@ -102,20 +104,6 @@ def _pump_stdout(proc: subprocess.Popen, log_handle) -> None:
             log_handle.close()
 
 
-def _detect_device() -> str:
-    """cuda → mps → cpu, picked once at start time."""
-    try:
-        import torch
-
-        if torch.cuda.is_available():
-            return "cuda"
-        if torch.backends.mps.is_available():
-            return "mps"
-    except Exception:
-        pass
-    return "cpu"
-
-
 def _resolve_policy_path(policy_ref: str) -> str:
     """Turn a checkpoints API ref into a local path that lerobot accepts.
 
@@ -345,7 +333,7 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]:
             "lerobot.scripts.lerobot_rollout",
             "--strategy.type=base",
             f"--policy.path={policy_path}",
-            f"--policy.device={_detect_device()}",
+            f"--policy.device={auto_select_torch_device().type}",
             "--robot.type=so101_follower",
             f"--robot.port={request.follower_port}",
             f"--robot.id={follower_id}",

lelab/train.py

@@ -24,6 +24,8 @@
 
 from pydantic import BaseModel
 
+from lerobot.utils.device_utils import auto_select_torch_device
+
 if TYPE_CHECKING:
     from lelab.jobs import JobTarget
 
@@ -75,7 +77,7 @@ class TrainingRequest(BaseModel):
     eval_use_async_envs: bool = False
 
     # Policy-specific
-    policy_device: str | None = "cuda"
+    policy_device: str | None = None
     policy_use_amp: bool = False
     # Hub upload (set by HfCloudJobRunner; not exposed in the form)
     policy_push_to_hub: bool = False
@@ -145,15 +147,15 @@ def build_training_command(
         cmd.extend(["--seed", str(request.seed)])
 
     # Policy device / AMP / hub
-    if request.policy_device:
-        cmd.extend(["--policy.device", request.policy_device])
     cmd.extend(["--policy.use_amp", "true" if request.policy_use_amp else "false"])
     # On HF Cloud, lerobot's submit_to_hf owns the model repo and sets push_to_hub on
     # the pod itself; _pod_forwarded_args drops any --policy.push_to_hub/--policy.repo_id
     # we'd pass, so we must not emit them. Local runs keep the existing behavior:
     # LeRobot defaults push_to_hub=True and demands --policy.repo_id when so.
+    # submit_to_hf also resets policy.device so the pod picks its own device.
     is_cloud = job_target is not None and job_target.runner == "hf_cloud"
     if not is_cloud:
+        cmd.extend(["--policy.device", request.policy_device or auto_select_torch_device().type])
         cmd.extend(["--policy.push_to_hub", "true" if request.policy_push_to_hub else "false"])
         if request.policy_push_to_hub and request.policy_repo_id:
             cmd.extend(["--policy.repo_id", request.policy_repo_id])

tests/test_rollout.py

@@ -61,50 +61,6 @@ def test_inference_request_has_expected_defaults() -> None:
     assert req.duration_s == 60
 
 
-def test_detect_device_returns_cpu_when_neither_cuda_nor_mps(monkeypatch: pytest.MonkeyPatch) -> None:
-    import torch
-
-    from lelab.rollout import _detect_device
-
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
-    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False)
-    assert _detect_device() == "cpu"
-
-
-def test_detect_device_prefers_cuda_over_mps(monkeypatch: pytest.MonkeyPatch) -> None:
-    import torch
-
-    from lelab.rollout import _detect_device
-
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
-    assert _detect_device() == "cuda"
-
-
-def test_detect_device_falls_back_to_mps_when_no_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
-    import torch
-
-    from lelab.rollout import _detect_device
-
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
-    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
-    assert _detect_device() == "mps"
-
-
-def test_detect_device_returns_cpu_when_torch_probe_raises(monkeypatch: pytest.MonkeyPatch) -> None:
-    """The function wraps both probes in a broad try/except — if torch is
-    broken at runtime we still need a sensible fallback."""
-    import torch
-
-    from lelab.rollout import _detect_device
-
-    def _boom() -> bool:
-        raise RuntimeError("simulated torch.cuda failure")
-
-    monkeypatch.setattr(torch.cuda, "is_available", _boom)
-    assert _detect_device() == "cpu"
-
-
 def test_resolve_policy_path_returns_local_dir_unchanged(tmp_path) -> None:
     from lelab.rollout import _resolve_policy_path
 

tests/test_train.py

@@ -137,14 +137,18 @@ def test_cloud_target_emits_job_flags_and_skips_push_to_hub() -> None:
 
     # push_to_hub is requested, but for a cloud target lerobot's submit_to_hf owns the
     # model repo and _pod_forwarded_args drops --policy.* — so we must NOT emit them.
-    req = TrainingRequest(dataset_repo_id="x", policy_push_to_hub=True, policy_repo_id="me/x")
+    req = TrainingRequest(
+        dataset_repo_id="x", policy_push_to_hub=True, policy_repo_id="me/x", policy_device="mps"
+    )
     cmd = build_training_command(
         req, "/tmp/out", job_target=JobTarget(runner="hf_cloud", flavor="a10g-small")
     )
     assert _arg_value(cmd, "--job.target") == "a10g-small"
     assert _arg_value(cmd, "--job.tags") == '["lelab"]'
     assert "--policy.push_to_hub" not in cmd
     assert "--policy.repo_id" not in cmd
+    # submit_to_hf resets policy.device so the pod picks its own device.
+    assert "--policy.device" not in cmd
     # An absolute host output_dir would be baked into the staged config and crash the
     # pod (mkdir /Users ...); checkpoints go to the Hub repo, so it must be omitted.
     assert "--output_dir" not in cmd
@@ -177,6 +181,23 @@ def test_local_target_keeps_push_to_hub() -> None:
     assert "--job.target" not in cmd
 
 
+def test_local_run_uses_detected_device_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    import torch
+
+    import lelab.train as train
+
+    monkeypatch.setattr(train, "auto_select_torch_device", lambda: torch.device("mps"))
+    cmd = train.build_training_command(train.TrainingRequest(dataset_repo_id="x"), "/tmp/out")
+    assert _arg_value(cmd, "--policy.device") == "mps"
+
+
+def test_local_run_keeps_explicit_device() -> None:
+    from lelab.train import TrainingRequest, build_training_command
+
+    cmd = build_training_command(TrainingRequest(dataset_repo_id="x", policy_device="cpu"), "/tmp/out")
+    assert _arg_value(cmd, "--policy.device") == "cpu"
+
+
 def test_groot_policy_flags_are_emitted_only_for_groot() -> None:
     from lelab.train import TrainingRequest, build_training_command