refactor: reshape tensor in 4d, then pad

Author: Federico Pozzi

test/test_prototype_transforms_functional.py

@@ -377,12 +377,6 @@ def pad_segmentation_mask():
         [[1], [1, 1], [1, 1, 2, 2]],  # padding
         ["constant", "symmetric", "edge", "reflect"],  # padding mode,
     ):
-        if padding_mode == "symmetric" and mask.ndim not in [3, 4]:
-            continue
-
-        if (padding_mode == "edge" or padding_mode == "reflect") and mask.ndim not in [3, 4]:
-            continue
-
         yield SampleInput(mask, padding=padding, padding_mode=padding_mode)
 
 

torchvision/prototype/transforms/functional/_geometry.py

@@ -399,7 +399,15 @@ def rotate_segmentation_mask(
 def pad_segmentation_mask(
     segmentation_mask: torch.Tensor, padding: List[int], padding_mode: str = "constant"
 ) -> torch.Tensor:
-    return pad_image_tensor(img=segmentation_mask, padding=padding, fill=0, padding_mode=padding_mode)
+    num_masks, height, width = segmentation_mask.shape[-3:]
+    extra_dims = segmentation_mask.shape[:-3]
+
+    padded_mask = pad_image_tensor(
+        img=segmentation_mask.view(-1, num_masks, height, width), padding=padding, fill=0, padding_mode=padding_mode
+    )
+
+    new_height, new_width = padded_mask.shape[-2:]
+    return padded_mask.view(extra_dims + (num_masks, new_height, new_width))
 
 
 def pad_bounding_box(