[v2] Preserve channels_last memory format in pad (#9560)

test/test_transforms_v2.py

@@ -4860,6 +4860,28 @@ def test_image_correctness(self, padding, padding_mode, fill, fn):
 
         assert_equal(actual, expected)
 
+    @pytest.mark.parametrize(
+        ("padding_mode", "fill"),
+        [
+            ("constant", 7),
+            ("edge", None),
+            ("reflect", None),
+            ("symmetric", None),
+            ("constant", [1, 2, 3]),
+        ],
+    )
+    @pytest.mark.parametrize("fn", [F.pad, F.pad_image])
+    def test_channels_last_memory_format(self, padding_mode, fill, fn):
+        image = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5)
+        image = image.to(memory_format=torch.channels_last)
+        padding = [1, 1, 1, 1]
+
+        actual = fn(image, padding=padding, fill=fill, padding_mode=padding_mode)
+        expected = F.pad(image.contiguous(), padding=padding, fill=fill, padding_mode=padding_mode)
+
+        torch.testing.assert_close(actual, expected)
+        assert actual.is_contiguous(memory_format=torch.channels_last)
+
     def _reference_pad_bounding_boxes(self, bounding_boxes, *, padding):
         if isinstance(padding, int):
             padding = [padding]

torchvision/transforms/v2/functional/_geometry.py

@@ -1628,6 +1628,7 @@ def _pad_with_scalar_fill(
         batch_size *= s
 
     image = image.reshape(batch_size, num_channels, height, width)
+    preserve_channels_last = image.is_contiguous(memory_format=torch.channels_last) and not image.is_contiguous()
 
     if padding_mode == "edge":
         # Similar to the padding order, `torch_pad`'s PIL's padding modes don't have the same names. Thus, we map
@@ -1656,6 +1657,9 @@ def _pad_with_scalar_fill(
 
     new_height, new_width = image.shape[-2:]
 
+    if preserve_channels_last:
+        image = image.contiguous(memory_format=torch.channels_last)
+
     return image.reshape(shape[:-3] + (num_channels, new_height, new_width))