From 6589a1b51a1d758ff99df4a2936c3b79786db113 Mon Sep 17 00:00:00 2001 From: 96sanjay Date: Fri, 1 May 2026 17:37:52 +0200 Subject: [PATCH] fix(saute): fix device mismatch in SauteAdapter with custom envs When the environment returns CPU tensors while training on GPU, SauteAdapter crashes on the first step because _safety_obs lives on the training device but env outputs (obs, reward, cost, etc.) are still on CPU. Coerce all env outputs to self._device at the step() and reset() boundary. Also guard _augment_obs for final_observation from the info dict. No-op when tensors are already on the correct device. --- omnisafe/adapter/saute_adapter.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/omnisafe/adapter/saute_adapter.py b/omnisafe/adapter/saute_adapter.py index 95a136605..7c25c0389 100644 --- a/omnisafe/adapter/saute_adapter.py +++ b/omnisafe/adapter/saute_adapter.py @@ -128,6 +128,7 @@ def reset( info: Some information logged by the environment. """ obs, info = self._env.reset(seed=seed, options=options) + obs = obs.to(self._device) self._safety_obs = torch.ones(self._env.num_envs, 1).to(self._device) obs = self._augment_obs(obs) return obs, info @@ -161,6 +162,11 @@ def step( info: Some information logged by the environment. """ next_obs, reward, cost, terminated, truncated, info = self._env.step(action) + next_obs = next_obs.to(self._device) + reward = reward.to(self._device) + cost = cost.to(self._device) + terminated = terminated.to(self._device) + truncated = truncated.to(self._device) info['original_reward'] = reward self._safety_step(cost) @@ -214,6 +220,7 @@ def _augment_obs(self, obs: torch.Tensor) -> torch.Tensor: Returns: The augmented observation. """ + obs = obs.to(self._device) return torch.cat([obs, self._safety_obs], dim=-1) def _log_value(