Allow v2 Normalize to reuse tensor statistics

test/test_transforms_v2.py

@@ -5670,6 +5670,18 @@ def test_transform(self, make_input):
             check_sample_input=self._sample_input_adapter,
         )
 
+    def test_transform_accepts_tensor_statistics(self):
+        mean = torch.tensor(self.MEAN)
+        std = torch.tensor(self.STD)
+        transform = transforms.Normalize(mean=mean, std=std)
+
+        assert transform.mean is mean
+        assert transform.std is std
+        image = make_image(dtype=torch.float32)
+        actual = transform(image)
+        expected = F.normalize_image(image, mean=mean, std=std)
+        assert_equal(actual, expected)
+
     def _reference_normalize_image(self, image, *, mean, std):
         image = image.numpy()
         mean, std = (np.array(stat, dtype=image.dtype).reshape((-1, 1, 1)) for stat in [mean, std])

torchvision/transforms/v2/_misc.py

@@ -152,18 +152,19 @@ class Normalize(Transform):
         This transform acts out of place, i.e., it does not mutate the input tensor.
 
     Args:
-        mean (sequence): Sequence of means for each channel.
-        std (sequence): Sequence of standard deviations for each channel.
+        mean (sequence or Tensor): Sequence of means for each channel. A tensor can
+            be provided to keep the statistics on the same device as the input.
+        std (sequence or Tensor): Sequence of standard deviations for each channel.
         inplace(bool,optional): Bool to make this operation in-place.
 
     """
 
     _v1_transform_cls = _transforms.Normalize
 
-    def __init__(self, mean: Sequence[float], std: Sequence[float], inplace: bool = False):
+    def __init__(self, mean: Sequence[float] | torch.Tensor, std: Sequence[float] | torch.Tensor, inplace: bool = False):
         super().__init__()
-        self.mean = list(mean)
-        self.std = list(std)
+        self.mean = mean if isinstance(mean, torch.Tensor) else list(mean)
+        self.std = std if isinstance(std, torch.Tensor) else list(std)
         self.inplace = inplace
 
     def check_inputs(self, sample: Any) -> Any: