amend

pytorch · vmoens · May 31, 2024 · Apr 30, 2024 · May 26, 2024 · May 26, 2024
commit 4a502505dcd411a65373b0decc7e047a4934edd0
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -4918,5 +4918,7 @@ def _minmax_dtype(dtype):
     return info.min, info.max
 
 
-def _remove_neg_shapes(shape):
+def _remove_neg_shapes(*shape):
+    if len(shape) == 1 and not isinstance(shape[0], int):
+        return _remove_neg_shapes(*shape[0])
     return torch.Size([int(d) if d >= 0 else 1 for d in shape])
diff --git a/torchrl/data/utils.py b/torchrl/data/utils.py
@@ -75,7 +75,7 @@ def consolidate_spec(
     recurse_through_entries: bool = True,
     recurse_through_stack: bool = True,
 ):
-    """Given a TensorSpec, removes exclusive keys by adding 0 shaped specs.
+    """Given a TensorSpec, removes exclusive keys by adding -1 shaped specs.
 
     Args:
         spec (CompositeSpec): the spec to be consolidated.
@@ -173,7 +173,7 @@ def _empty_like_spec(specs: List[TensorSpec], shape):
             if hetero_dim:
                 spec_shape[dim_index] = -1
 
-        if 0 not in spec_shape:  # the values have all same shape
+        if -1 not in spec_shape:  # the values have all same shape
             spec_shape = [
                 dim if i < len(shape) else -1 for i, dim in enumerate(spec_shape)
             ]