Add expected shapes in tvtensors docs

torchvision/tv_tensors/_bounding_boxes.py

@@ -25,7 +25,7 @@ class BoundingBoxFormat(Enum):
 
 
 class BoundingBoxes(TVTensor):
-    """:class:`torch.Tensor` subclass for bounding boxes.
+    """:class:`torch.Tensor` subclass for bounding boxes with shape ``[N, 4]``.
 
     .. note::
         There should be only one :class:`~torchvision.tv_tensors.BoundingBoxes`

torchvision/tv_tensors/_image.py

@@ -9,7 +9,7 @@
 
 
 class Image(TVTensor):
-    """:class:`torch.Tensor` subclass for images.
+    """:class:`torch.Tensor` subclass for images with shape ``[..., C, H, W]``.
 
     .. note::
 

torchvision/tv_tensors/_mask.py

@@ -9,7 +9,7 @@
 
 
 class Mask(TVTensor):
-    """:class:`torch.Tensor` subclass for segmentation and detection masks.
+    """:class:`torch.Tensor` subclass for segmentation and detection masks with shape ``[..., H, W]``.
 
     Args:
         data (tensor-like, PIL.Image.Image): Any data that can be turned into a tensor with :func:`torch.as_tensor` as

torchvision/tv_tensors/_video.py

@@ -8,7 +8,7 @@
 
 
 class Video(TVTensor):
-    """:class:`torch.Tensor` subclass for videos.
+    """:class:`torch.Tensor` subclass for videos with shape ``[..., T, C, H, W]``.
 
     Args:
         data (tensor-like): Any data that can be turned into a tensor with :func:`torch.as_tensor`.