How to use the torchgeometry.color.normalize function in torchgeometry

def op_script(data: torch.Tensor, mean: torch.Tensor,
                      std: torch.Tensor) -> torch.Tensor:

            return color.normalize(data, mean, std)

            data = torch.ones(2, 3, 1, 1)
            data += 2

            mean = torch.tensor([0.5, 1.0, 2.0]).repeat(2, 1)
            std = torch.tensor([2.0, 2.0, 2.0]).repeat(2, 1)

            actual = op_script(data, mean, std)
            expected = color.normalize(data, mean, std)
            assert_allclose(actual, expected)

def op_script(data: torch.Tensor, mean: torch.Tensor,
                      std: torch.Tensor) -> torch.Tensor:

            return color.normalize(data, mean, std)

            data = torch.ones(2, 3, 1, 1)
            data += 2

            mean = torch.tensor([0.5, 1.0, 2.0]).repeat(2, 1)
            std = torch.tensor([2.0, 2.0, 2.0]).repeat(2, 1)

            actual = op_script(data, mean, std)
            expected = color.normalize(data, mean, std)
            assert_allclose(actual, expected)