Add shape to BatchedTensor (#237)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redcat"
-version = "0.0.1a193"
+version = "0.0.1a194"
 description = "A library to manipulate batches of examples"
 readme = "README.md"
 authors = ["Thibaut Durand <durand.tibo+gh@gmail.com>"]

src/redcat/tensor.py

@@ -80,6 +80,11 @@ def device(self) -> torch.device:
         r"""``torch.device``: The device where the batch data/tensor is."""
         return self._data.device
 
+    @property
+    def shape(self) -> torch.Size:
+        r"""``torch.Size``: The shape of the tensor."""
+        return self._data.shape
+
     def dim(self) -> int:
         r"""Gets the number of dimensions.
 

tests/unit/test_tensor.py

@@ -70,6 +70,10 @@ def test_batched_tensor_device(device: str) -> None:
     assert BatchedTensor(torch.ones(2, 3, device=device)).device == device
 
 
+def test_batched_tensor_shape() -> None:
+    assert BatchedTensor(torch.ones(2, 3)).shape == torch.Size([2, 3])
+
+
 def test_batched_tensor_dim() -> None:
     assert BatchedTensor(torch.ones(2, 3)).dim() == 2
 

tests/unit/test_tensorseq.py

@@ -108,6 +108,10 @@ def test_batched_tensor_seq_device(device: str) -> None:
     assert BatchedTensorSeq(torch.ones(2, 3, device=device)).device == device
 
 
+def test_batched_tensor_seq_shape() -> None:
+    assert BatchedTensorSeq(torch.ones(2, 3)).shape == torch.Size([2, 3])
+
+
 def test_batched_tensor_seq_dim() -> None:
     assert BatchedTensorSeq(torch.ones(2, 3)).dim() == 2