improve UX for v2 Compose

test/test_transforms_v2_refactored.py

@@ -26,6 +26,8 @@
     make_video,
     set_rng_seed,
 )
+
+from torch import nn
 from torch.testing import assert_close
 from torchvision import datapoints
 
@@ -1634,3 +1636,56 @@ def test_transform_negative_degrees_error(self):
     def test_transform_unknown_fill_error(self):
         with pytest.raises(TypeError, match="Got inappropriate fill arg"):
             transforms.RandomAffine(degrees=0, fill="fill")
+
+
+class TestCompose:
+    class BuiltinTransform(transforms.Transform):
+        def _transform(self, inpt, params):
+            return inpt
+
+    class PackedInputTransform(nn.Module):
+        def forward(self, sample):
+            image, label = sample
+            return image, label
+
+    class UnpackedInputTransform(nn.Module):
+        def forward(self, image, label):
+            return image, label
+
+    @pytest.mark.parametrize(
+        "transform_clss",
+        [
+            [BuiltinTransform],
+            [PackedInputTransform],
+            [UnpackedInputTransform],
+            [BuiltinTransform, BuiltinTransform],
+            [PackedInputTransform, PackedInputTransform],
+            [UnpackedInputTransform, UnpackedInputTransform],
+            [BuiltinTransform, PackedInputTransform, BuiltinTransform],
+            [BuiltinTransform, UnpackedInputTransform, BuiltinTransform],
+            [PackedInputTransform, BuiltinTransform, PackedInputTransform],
+            [UnpackedInputTransform, BuiltinTransform, UnpackedInputTransform],
+        ],
+    )
+    @pytest.mark.parametrize("unpack", [True, False])
+    def test_packed_unpacked(self, transform_clss, unpack):
+        if any(
+            unpack
+            and issubclass(cls, self.PackedInputTransform)
+            or not unpack
+            and issubclass(cls, self.UnpackedInputTransform)
+            for cls in transform_clss
+        ):
+            return
+
+        transform = transforms.Compose([cls() for cls in transform_clss])
+
+        image = make_image()
+        label = torch.tensor(3)
+        packed_input = (image, label)
+
+        output = transform(*packed_input if unpack else (packed_input,))
+
+        assert isinstance(output, tuple) and len(output) == 2
+        assert output[0] is image
+        assert output[1] is label

torchvision/transforms/v2/_container.py

@@ -46,10 +46,15 @@ def __init__(self, transforms: Sequence[Callable]) -> None:
         self.transforms = transforms
 
     def forward(self, *inputs: Any) -> Any:
-        sample = inputs if len(inputs) > 1 else inputs[0]
+        needs_unpacking = len(inputs) > 1
+
+        if not self.transforms:
+            return inputs if needs_unpacking else inputs[0]
+
         for transform in self.transforms:
-            sample = transform(sample)
-        return sample
+            outputs = transform(*inputs)
+            inputs = outputs if needs_unpacking else (outputs,)
+        return outputs
 
     def extra_repr(self) -> str:
         format_string = []