[BugFix] better device consistency in EGreedy (#1867)

pytorch · Feb 4, 2024 · 5f82601 · 5f82601
1 parent 0672359
commit 5f82601
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/torchrl/modules/tensordict_module/exploration.py b/torchrl/modules/tensordict_module/exploration.py
@@ -149,10 +149,7 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
 
             out = action_tensordict.get(action_key)
             eps = self.eps.item()
-            cond = (
-                torch.rand(action_tensordict.shape, device=action_tensordict.device)
-                < eps
-            ).to(out.dtype)
+            cond = torch.rand(action_tensordict.shape, device=out.device) < eps
             cond = expand_as_right(cond, out)
             spec = self.spec
             if spec is not None:
@@ -177,7 +174,7 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
                             f"Action mask key {self.action_mask_key} not found in {tensordict}."
                         )
                     spec.update_mask(action_mask)
-                out = cond * spec.rand().to(out.device) + (1 - cond) * out
+                out = torch.where(cond, spec.rand().to(out.device), out)
             else:
                 raise RuntimeError("spec must be provided to the exploration wrapper.")
             action_tensordict.set(action_key, out)