Update

[ghstack-poisoned]
pytorch · Dec 14, 2024 · ed89a10 · ed89a10
1 parent 3da531a
commit ed89a10
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -1695,7 +1695,7 @@ def unbind(self, dim: int = 0):
             for i in range(self.shape[dim])
         )
 
-    @implement_for("torch", None, "2.1")
+    @implement_for("torch", None, "2.1", compilable=True)
     def rand(self, shape: torch.Size = None) -> torch.Tensor:
         if shape is None:
             shape = self.shape[:-1]
@@ -1718,7 +1718,7 @@ def rand(self, shape: torch.Size = None) -> torch.Tensor:
         # out.scatter_(-1, m, 1)
         return out
 
-    @implement_for("torch", "2.1")
+    @implement_for("torch", "2.1", compilable=True)
     def rand(self, shape: torch.Size = None) -> torch.Tensor:  # noqa: F811
         if shape is None:
             shape = self.shape[:-1]