diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index 6d9f9732552..be60ac90357 100644 --- a/test/test_transforms_v2.py +++ b/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] diff --git a/torchvision/transforms/v2/functional/_geometry.py b/torchvision/transforms/v2/functional/_geometry.py index 3ebf208b1a1..388c27e65a9 100644 --- a/torchvision/transforms/v2/functional/_geometry.py +++ b/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))