Added tests on batch of tensors to check transforms

docs/source/transforms.rst

@@ -9,6 +9,11 @@ Functional transforms give fine-grained control over the transformations.
 This is useful if you have to build a more complex transformation pipeline
 (e.g. in the case of segmentation tasks).
 
+All transformations accept PIL Image, Tensor Image or batch of Tensor Images as input. Tensor Image is a tensor with
+``(C, H, W)`` shape, where ``C`` is a number of channels, ``H`` and ``W`` are image height and width. Batch of
+Tensor Images is a tensor of ``(B, C, H, W)`` shape, where ``B`` is a number of images in the batch. Deterministic or
+random transformations applied on the batch of Tensor Images identically transform all the images of the batch.
+
 .. autoclass:: Compose
 
 Transforms on PIL Image

test/common_utils.py

@@ -341,6 +341,15 @@ def _create_data(self, height=3, width=3, channels=3, device="cpu"):
         pil_img = Image.fromarray(tensor.permute(1, 2, 0).contiguous().cpu().numpy())
         return tensor, pil_img
 
+    def _create_data_batch(self, height=3, width=3, channels=3, num_samples=4, device="cpu"):
+        batch_tensor = torch.randint(
+            0, 255,
+            (num_samples, channels, height, width),
+            dtype=torch.uint8,
+            device=device
+        )
+        return batch_tensor
+
     def compareTensorToPIL(self, tensor, pil_image, msg=None):
         np_pil_image = np.array(pil_image)
         if np_pil_image.ndim == 2: