How to use the kornia.color function in kornia

def test_jit(self):
        batch_size, channels, height, width = 2, 3, 64, 64
        img = torch.ones(batch_size, channels, height, width)
        gray = kornia.color.RgbToGrayscale()
        gray_traced = torch.jit.trace(kornia.color.RgbToGrayscale(), img)
        assert_allclose(gray(img), gray_traced(img))

def test_jit(self):
        batch_size, channels, height, width = 2, 3, 64, 64
        img = torch.ones(batch_size, channels, height, width)
        gray = kornia.color.RgbToGrayscale()
        gray_traced = torch.jit.trace(kornia.color.RgbToGrayscale(), img)
        assert_allclose(gray(img), gray_traced(img))

def test_gamma_one(self):
        data = torch.tensor([[[1., 1.],
                              [1., 1.]],

                             [[.5, .5],
                              [.5, .5]],

                             [[.25, .25],
                              [.25, .25]]])  # 3x2x2

        expected = data
        f = kornia.color.AdjustGamma(1.)
        assert_allclose(f(data), expected)

def test_smoke(self):
        mean = [0.5]
        std = [0.1]
        repr = 'Normalize(mean=[0.5], std=[0.1])'
        assert str(kornia.color.Normalize(mean, std)) == repr

def test_saturation_one(self):
        data = torch.tensor([[[.5, .5],
                              [.5, .5]],

                             [[.5, .5],
                              [.5, .5]],

                             [[.25, .25],
                              [.25, .25]]])  # 3x2x2

        expected = data
        f = kornia.color.AdjustSaturation(1.)
        assert_allclose(f(data), expected)

def test_gradcheck(self, size, device):
        shape = random_shape(size, max_elem=5)  # to shave time on gradcheck
        src1 = torch.randn(shape).to(device)
        src2 = torch.randn(shape).to(device)
        alpha = random.random()
        beta = random.random()
        gamma = random.random()

        src1 = utils.tensor_to_gradcheck_var(src1)  # to var
        src2 = utils.tensor_to_gradcheck_var(src2)  # to var

        assert gradcheck(kornia.color.AddWeighted(alpha, beta, gamma), (src1, src2),
                         raise_exception=True)

def test_normalize(self):

        # prepare input data
        data = torch.ones(1, 2, 2)
        mean = torch.tensor([0.5])
        std = torch.tensor([2.0])

        # expected output
        expected = torch.tensor([0.25]).repeat(1, 2, 2).view_as(data)

        f = kornia.color.Normalize(mean, std)
        assert_allclose(f(data), expected)

[.5, .5]],

                              [[.25, .25],
                               [.25, .25]]],

                             [[[.5, .5],
                               [.5, .5]],

                              [[.5, .5],
                               [.5, .5]],

                              [[.25, .25],
                               [.25, .25]]]])  # 2x3x2x2

        expected = data
        f = kornia.color.AdjustHue(torch.tensor([0, 0]))
        assert_allclose(f(data), expected)

def test_addweighted(self, size, device):
        src1, src2, alpha, beta, gamma = self.get_input(3)
        src1 = src1.to(device)
        src2 = src2.to(device)

        f = kornia.color.AddWeighted(alpha, beta, gamma)
        assert_allclose(f(src1, src2), src1 * alpha + src2 * beta + gamma)

[3., 3.]]])  # 3x2x2

        expected = torch.tensor([[[3., 3.],
                                  [3., 3.]],

                                 [[2., 2.],
                                  [2., 2.]],

                                 [[1., 1.],
                                  [1., 1.]]])  # 3x2x2

        # move data to the device
        data = data.to(device)
        expected = expected.to(device)

        f = kornia.color.BgrToRgb()
        assert_allclose(f(data), expected)