[RLlib] Fix: Added dtype to torch deterministic action. (ray-project#…

…29648) Signed-off-by: Weichen Xu <[email protected]>
WeichenXu123 · Dec 19, 2022 · e2388c4 · e2388c4
1 parent 67625d6
commit e2388c4
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/rllib/models/torch/torch_distributions.py b/rllib/models/torch/torch_distributions.py
@@ -219,8 +219,11 @@ def sample(
 
         if sample_shape is None:
             sample_shape = torch.Size()
-        loc_shape = self.loc.shape
-        return torch.ones(sample_shape + loc_shape, device=self.loc.device) * self.loc
+
+        device = self.loc.device
+        dtype = self.loc.dtype
+        shape = sample_shape + self.loc.shape
+        return torch.ones(shape, device=device, dtype=dtype) * self.loc
 
     def rsample(
         self,