diff --git a/source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst b/source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
new file mode 100644
index 000000000000..ba2c7b78ec0d
--- /dev/null
+++ b/source/isaaclab/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed PyTorch and Warp selecting different CUDA devices during simulation
+  initialization.
diff --git a/source/isaaclab/isaaclab/app/app_launcher.py b/source/isaaclab/isaaclab/app/app_launcher.py
index b330724f39e9..d538e3b68024 100644
--- a/source/isaaclab/isaaclab/app/app_launcher.py
+++ b/source/isaaclab/isaaclab/app/app_launcher.py
@@ -29,6 +29,7 @@
     from isaacsim import SimulationApp
 
 from isaaclab.app.settings_manager import get_settings_manager, initialize_carb_settings
+from isaaclab.utils._device import set_cuda_device
 
 # import logger
 logger = logging.getLogger(__name__)
@@ -1081,13 +1082,11 @@ def _resolve_device_settings(self, launcher_args: dict):
         logger.info("Using device: %s", device)
 
     def _set_deferred_cuda_device(self) -> None:
-        """Set the current torch CUDA device after Kit startup."""
+        """Set the current CUDA device after Kit startup."""
         if self._deferred_cuda_device_id is None:
             return
 
-        import torch
-
-        torch.cuda.set_device(self._deferred_cuda_device_id)
+        set_cuda_device(self._deferred_cuda_device_id)
 
     def _resolve_experience_file(self, launcher_args: dict):
         """Resolve experience file related settings."""
diff --git a/source/isaaclab/isaaclab/app/sim_launcher.py b/source/isaaclab/isaaclab/app/sim_launcher.py
index fc69dd3cdc97..c1f82afc1dd9 100644
--- a/source/isaaclab/isaaclab/app/sim_launcher.py
+++ b/source/isaaclab/isaaclab/app/sim_launcher.py
@@ -30,6 +30,7 @@
 from isaaclab.physics.physics_manager_cfg import PhysicsCfg
 from isaaclab.renderers.renderer_cfg import RendererCfg
 from isaaclab.sensors.camera.camera_cfg import CameraCfg
+from isaaclab.utils._device import set_cuda_device
 
 logger = logging.getLogger(__name__)
 
@@ -379,7 +380,7 @@ def _resolve_distributed_device(cfg, launcher_args: argparse.Namespace | dict |
     sim_cfg = getattr(cfg, "sim", None)
     if sim_cfg is not None:
         sim_cfg.device = device_str
-    torch.cuda.set_device(device_str)
+    set_cuda_device(device_str)
     logger.info(
         "Distributed device resolved to %s (local_rank=%d, visible_gpus=%d)",
         device_str,
diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py
index 5255af873604..5c1a5d4403f1 100644
--- a/source/isaaclab/isaaclab/physics/physics_manager.py
+++ b/source/isaaclab/isaaclab/physics/physics_manager.py
@@ -14,6 +14,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, ClassVar
 
+from isaaclab.utils._device import set_cuda_device
+
 if TYPE_CHECKING:
     from isaaclab.scene_data import SceneDataBackend
     from isaaclab.sim.simulation_context import SimulationContext
@@ -271,6 +273,12 @@ def initialize(cls, sim_context: SimulationContext) -> None:
         PhysicsManager._device = sim_context.cfg.device
         PhysicsManager._sim_time = 0.0
 
+        # Synchronize the process-wide CUDA device before backend-specific
+        # initialization allocates state. PyTorch must select the device before
+        # Warp so that both runtimes retain the same primary CUDA context.
+        if "cuda" in PhysicsManager._device:
+            set_cuda_device(PhysicsManager._device)
+
     @classmethod
     @abstractmethod
     def reset(cls, soft: bool = False) -> None:
diff --git a/source/isaaclab/isaaclab/utils/_device.py b/source/isaaclab/isaaclab/utils/_device.py
new file mode 100644
index 000000000000..020be59dfec2
--- /dev/null
+++ b/source/isaaclab/isaaclab/utils/_device.py
@@ -0,0 +1,27 @@
+# Copyright (c) 2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Internal utilities for selecting compute devices."""
+
+from __future__ import annotations
+
+
+def set_cuda_device(device: str | int) -> None:
+    """Set the process-wide CUDA device for both PyTorch and Warp.
+
+    PyTorch must select the device before Warp so that a newly imported Warp
+    runtime initializes on that device, or an existing runtime switches to it.
+
+    Args:
+        device: CUDA device index or a device string such as ``"cuda:1"``.
+    """
+    import torch
+
+    torch.cuda.set_device(device)
+
+    import warp as wp
+
+    warp_device = f"cuda:{device}" if isinstance(device, int) else device
+    wp.set_device(warp_device)
diff --git a/source/isaaclab/test/app/test_kwarg_launch.py b/source/isaaclab/test/app/test_kwarg_launch.py
index 1e3782840a8f..10aa804f973f 100644
--- a/source/isaaclab/test/app/test_kwarg_launch.py
+++ b/source/isaaclab/test/app/test_kwarg_launch.py
@@ -80,6 +80,18 @@ def __init__(self, _launcher_args):
     assert close_args == {"exit_code": 1}
 
 
+def test_deferred_cuda_device_synchronizes_torch_and_warp(monkeypatch: pytest.MonkeyPatch):
+    """The post-Kit device hook must synchronize both CUDA runtimes."""
+    devices = []
+    monkeypatch.setattr(app_launcher_module, "set_cuda_device", devices.append)
+    launcher = AppLauncher.__new__(AppLauncher)
+    launcher._deferred_cuda_device_id = 2
+
+    launcher._set_deferred_cuda_device()
+
+    assert devices == [2]
+
+
 class _DummySettings:
     def __init__(self):
         self.values = {}
diff --git a/source/isaaclab/test/devices/test_physics_manager_device.py b/source/isaaclab/test/devices/test_physics_manager_device.py
new file mode 100644
index 000000000000..6b8face7f304
--- /dev/null
+++ b/source/isaaclab/test/devices/test_physics_manager_device.py
@@ -0,0 +1,46 @@
+# Copyright (c) 2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for physics-manager CUDA device initialization."""
+
+from types import SimpleNamespace
+
+from isaaclab.physics import PhysicsManager
+from isaaclab.physics import physics_manager as physics_manager_module
+
+
+def test_initialize_synchronizes_cuda_before_backend_setup(monkeypatch):
+    """The base manager must synchronize CUDA before backend-specific setup."""
+    calls = []
+    monkeypatch.setattr(physics_manager_module, "set_cuda_device", lambda device: calls.append(("cuda", device)))
+    monkeypatch.setattr(PhysicsManager, "_sim", None)
+    monkeypatch.setattr(PhysicsManager, "_cfg", None)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:0")
+
+    class _TestManager(PhysicsManager):
+        @classmethod
+        def initialize(cls, sim_context):
+            super().initialize(sim_context)
+            calls.append(("backend", PhysicsManager._device))
+
+    sim_context = SimpleNamespace(cfg=SimpleNamespace(physics=object(), device="cuda:2"))
+
+    _TestManager.initialize(sim_context)
+
+    assert calls == [("cuda", "cuda:2"), ("backend", "cuda:2")]
+
+
+def test_initialize_does_not_synchronize_cpu_device(monkeypatch):
+    """CPU simulation must not initialize or select CUDA runtimes."""
+    devices = []
+    monkeypatch.setattr(physics_manager_module, "set_cuda_device", devices.append)
+    monkeypatch.setattr(PhysicsManager, "_sim", None)
+    monkeypatch.setattr(PhysicsManager, "_cfg", None)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:0")
+    sim_context = SimpleNamespace(cfg=SimpleNamespace(physics=object(), device="cpu"))
+
+    PhysicsManager.initialize(sim_context)
+
+    assert devices == []
diff --git a/source/isaaclab/test/utils/test_device.py b/source/isaaclab/test/utils/test_device.py
new file mode 100644
index 000000000000..7bfeacec1339
--- /dev/null
+++ b/source/isaaclab/test/utils/test_device.py
@@ -0,0 +1,37 @@
+# Copyright (c) 2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for compute-device selection utilities."""
+
+import sys
+from types import SimpleNamespace
+
+from isaaclab.utils._device import set_cuda_device
+
+
+def test_set_cuda_device_sets_torch_before_warp(monkeypatch):
+    """Setting a CUDA device must update PyTorch before Warp."""
+    calls = []
+    torch = SimpleNamespace(cuda=SimpleNamespace(set_device=lambda device: calls.append(("torch", device))))
+    warp = SimpleNamespace(set_device=lambda device: calls.append(("warp", device)))
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.setitem(sys.modules, "warp", warp)
+
+    set_cuda_device("cuda:2")
+
+    assert calls == [("torch", "cuda:2"), ("warp", "cuda:2")]
+
+
+def test_set_cuda_device_normalizes_integer_for_warp(monkeypatch):
+    """An integer device index must be converted to a Warp CUDA alias."""
+    calls = []
+    torch = SimpleNamespace(cuda=SimpleNamespace(set_device=lambda device: calls.append(("torch", device))))
+    warp = SimpleNamespace(set_device=lambda device: calls.append(("warp", device)))
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.setitem(sys.modules, "warp", warp)
+
+    set_cuda_device(3)
+
+    assert calls == [("torch", 3), ("warp", "cuda:3")]
diff --git a/source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst b/source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
new file mode 100644
index 000000000000..c53da690a8d8
--- /dev/null
+++ b/source/isaaclab_newton/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.rst
@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed Newton physics failing to initialize on non-default CUDA devices
+  (``cuda:1`` and higher).
diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
index 69fdc6fe1d07..82d412eae055 100644
--- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
+++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
@@ -1536,7 +1536,7 @@ def _capture_or_defer_graph(cls) -> None:
             with Timer(name="newton_cuda_graph", msg="CUDA graph took:"):
                 if cls._usdrt_stage is None:
                     simulate = cls._simulate_full if cls._is_all_graphable() else cls._simulate_physics_only
-                    with wp.ScopedCapture() as capture:
+                    with wp.ScopedCapture(device=device) as capture:
                         simulate()
                     NewtonManager._graph = capture.graph
                     logger.info("Newton CUDA graph captured (standard Warp mode)")
diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py
index 56bf4c33f1f1..f47e3a3232f4 100644
--- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py
+++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py
@@ -426,6 +426,38 @@ def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch):
     assert NewtonManager._graph_capture_pending is False
 
 
+def test_cuda_graph_capture_uses_simulation_device(monkeypatch):
+    """CUDA graph capture should use the simulation device instead of Warp's default device."""
+    from isaaclab.physics import PhysicsManager
+
+    captured_devices = []
+    captured_graph = object()
+
+    class FakeScopedCapture:
+        def __init__(self, device=None):
+            captured_devices.append(device)
+            self.graph = captured_graph
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            return False
+
+    monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False)
+    monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_solver", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False))
+    monkeypatch.setattr(NewtonManager, "_simulate_physics_only", classmethod(lambda cls: None))
+    monkeypatch.setattr(wp, "ScopedCapture", FakeScopedCapture)
+
+    NewtonManager._capture_or_defer_graph()
+
+    assert captured_devices == ["cuda:1"]
+    assert NewtonManager._graph is captured_graph
+
+
 # ---------------------------------------------------------------------------
 # Manager class hierarchy and factory contracts
 # ---------------------------------------------------------------------------
diff --git a/source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip b/source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip
new file mode 100644
index 000000000000..f03e45edd95d
--- /dev/null
+++ b/source/isaaclab_tasks/changelog.d/jichuanh-fix-newton-mgpu-capture-device-codex.skip
@@ -0,0 +1 @@
+Test-only updates for shared CUDA device-selection coverage.
diff --git a/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py b/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py
index 67268cb816c4..33de34289b59 100644
--- a/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py
+++ b/source/isaaclab_tasks/test/core/test_distributed_device_resolution.py
@@ -15,8 +15,8 @@
 - Non-distributed: no device override applied
 - launch_simulation device propagation from AppLauncher
 
-No actual GPUs required — ``torch.cuda.device_count`` and
-``torch.cuda.set_device`` are mocked throughout.
+No actual GPUs required — ``torch.cuda.device_count`` and the shared CUDA
+device-selection helper are mocked throughout.
 """
 
 from __future__ import annotations
@@ -95,7 +95,7 @@ def _make_env_vars(
 class TestResolveDistributedDeviceNamespace:
     """Tests for _resolve_distributed_device with argparse.Namespace args."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_normal_multi_gpu_rank0(self, mock_count, mock_set_device):
         """4 visible GPUs, world_size=4, rank 0 → cuda:0."""
@@ -109,7 +109,7 @@ def test_normal_multi_gpu_rank0(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_normal_multi_gpu_rank3(self, mock_count, mock_set_device):
         """4 visible GPUs, world_size=4, rank 3 → cuda:3."""
@@ -123,7 +123,7 @@ def test_normal_multi_gpu_rank3(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:3"
         mock_set_device.assert_called_once_with("cuda:3")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_cuda_visible_devices_restricted_rank0(self, mock_count, mock_set_device):
         """1 visible GPU (CUDA_VISIBLE_DEVICES set), world_size=2, rank 0 → cuda:0."""
@@ -137,7 +137,7 @@ def test_cuda_visible_devices_restricted_rank0(self, mock_count, mock_set_device
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_cuda_visible_devices_restricted_rank1(self, mock_count, mock_set_device):
         """1 visible GPU, world_size=2, rank 1 → falls back to cuda:0 (not cuda:1)."""
@@ -151,7 +151,7 @@ def test_cuda_visible_devices_restricted_rank1(self, mock_count, mock_set_device
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=2)
     def test_jax_local_rank_added(self, mock_count, mock_set_device):
         """JAX_LOCAL_RANK is added to LOCAL_RANK for correct device mapping."""
@@ -166,7 +166,7 @@ def test_jax_local_rank_added(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:1"
         mock_set_device.assert_called_once_with("cuda:1")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_jax_local_rank_with_restricted_gpus(self, mock_count, mock_set_device):
         """JAX_LOCAL_RANK + restricted GPUs → fallback to cuda:0."""
@@ -189,7 +189,7 @@ def test_jax_local_rank_with_restricted_gpus(self, mock_count, mock_set_device):
 class TestResolveDistributedDeviceDict:
     """Tests for _resolve_distributed_device with dict-style args."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_dict_args_distributed(self, mock_count, mock_set_device):
         """Dict launcher_args with distributed=True should work identically."""
@@ -203,7 +203,7 @@ def test_dict_args_distributed(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:2"
         mock_set_device.assert_called_once_with("cuda:2")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_dict_args_restricted(self, mock_count, mock_set_device):
         """Dict args with restricted GPUs should fall back to cuda:0."""
@@ -226,7 +226,7 @@ def test_dict_args_restricted(self, mock_count, mock_set_device):
 class TestResolveDistributedDeviceNoop:
     """Tests that non-distributed runs skip device resolution."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_not_distributed_namespace(self, mock_set_device):
         """distributed=False → device unchanged, set_device not called."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -237,7 +237,7 @@ def test_not_distributed_namespace(self, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_not_called()
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_not_distributed_dict(self, mock_set_device):
         """Dict with distributed=False → no-op."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -248,7 +248,7 @@ def test_not_distributed_dict(self, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_not_called()
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_no_distributed_key(self, mock_set_device):
         """Dict without 'distributed' key → no-op."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -259,7 +259,7 @@ def test_no_distributed_key(self, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_not_called()
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     def test_none_launcher_args(self, mock_set_device):
         """launcher_args=None → no-op."""
         env_cfg = _DummyEnvCfg(device="cuda:0")
@@ -278,7 +278,7 @@ def test_none_launcher_args(self, mock_set_device):
 class TestResolveDistributedDeviceEdgeCases:
     """Edge cases for device resolution."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=2)
     def test_env_cfg_without_sim(self, mock_count, mock_set_device):
         """env_cfg with no 'sim' attribute → set_device still called, no crash."""
@@ -296,7 +296,7 @@ class _BareEnvCfg:
         # Should still call set_device even without sim_cfg
         mock_set_device.assert_called_once_with("cuda:1")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=2)
     def test_world_size_equals_visible_gpus(self, mock_count, mock_set_device):
         """Exact match: 2 visible GPUs, world_size=2 → use local_rank directly."""
@@ -309,7 +309,7 @@ def test_world_size_equals_visible_gpus(self, mock_count, mock_set_device):
 
         assert env_cfg.sim.device == "cuda:1"
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=8)
     def test_more_gpus_than_world_size(self, mock_count, mock_set_device):
         """8 visible GPUs but only 2 ranks → use local_rank directly."""
@@ -322,7 +322,7 @@ def test_more_gpus_than_world_size(self, mock_count, mock_set_device):
 
         assert env_cfg.sim.device == "cuda:1"
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=0)
     def test_zero_visible_gpus(self, mock_count, mock_set_device):
         """0 visible GPUs → fallback to cuda:0 (will fail later at CUDA init)."""
@@ -335,7 +335,7 @@ def test_zero_visible_gpus(self, mock_count, mock_set_device):
 
         assert env_cfg.sim.device == "cuda:0"
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_missing_env_vars_default_to_zero(self, mock_count, mock_set_device):
         """Missing LOCAL_RANK/WORLD_SIZE → defaults to 0/1."""
@@ -364,7 +364,7 @@ def test_missing_env_vars_default_to_zero(self, mock_count, mock_set_device):
 class TestResolveDistributedDeviceMultiNode:
     """Tests for multi-node setups where WORLD_SIZE > local GPU count."""
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_multi_node_rank3_sees_4_gpus(self, mock_count, mock_set_device):
         """2 nodes × 4 GPUs, WORLD_SIZE=8, local_rank=3, 4 visible → cuda:3.
@@ -382,7 +382,7 @@ def test_multi_node_rank3_sees_4_gpus(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:3"
         mock_set_device.assert_called_once_with("cuda:3")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=4)
     def test_multi_node_rank0_sees_4_gpus(self, mock_count, mock_set_device):
         """2 nodes × 4 GPUs, WORLD_SIZE=8, local_rank=0 → cuda:0."""
@@ -396,7 +396,7 @@ def test_multi_node_rank0_sees_4_gpus(self, mock_count, mock_set_device):
         assert env_cfg.sim.device == "cuda:0"
         mock_set_device.assert_called_once_with("cuda:0")
 
-    @patch("torch.cuda.set_device")
+    @patch.object(sim_launcher, "set_cuda_device")
     @patch("torch.cuda.device_count", return_value=1)
     def test_multi_node_restricted_gpus(self, mock_count, mock_set_device):
         """Multi-node with CUDA_VISIBLE_DEVICES=<one GPU per rank>, local_rank=1 → cuda:0."""