[BugFix] Categorical spec samples the right dtype when masked (#2980)

Co-authored-by: Louis Faury <[email protected]>

Author: louisfaury

test/test_specs.py

@@ -1331,6 +1331,14 @@ def test_categorical_action_spec_rand(self):
         sample = [sum(sample == i) for i in range(10)]
         assert chisquare(sample).pvalue > 0.1
 
+    @pytest.mark.parametrize("dtype", [torch.int, torch.int32, torch.int64])
+    def test_categorical_action_spec_rand_masked_right_dtype(self, dtype: torch.dtype):
+        torch.manual_seed(1)
+        action_spec = Categorical(2, dtype=dtype)
+        action_spec.update_mask(torch.tensor([True, False]))
+        sample = action_spec.rand()
+        assert sample.dtype == dtype
+
     def test_mult_discrete_action_spec_rand(self):
         torch.manual_seed(0)
         ns = (10, 5)

torchrl/data/tensor_specs.py

@@ -3926,7 +3926,7 @@ def rand(self, shape: torch.Size = None) -> torch.Tensor:
                 "The last dimension of the mask must match the number of action allowed by the "
                 f"Categorical spec. Got mask.shape={self.mask.shape} and n={n}."
             )
-        out = torch.multinomial(mask_flat.float(), 1).reshape(shape_out)
+        out = torch.multinomial(mask_flat.float(), 1).reshape(shape_out).to(self.dtype)
         return out
 
     def index(