import torch

from kornia.testing import KORNIA_CHECK_IS_COLOR, KORNIA_CHECK_IS_TENSOR


def shift_rgb(image: torch.Tensor, r_shift: torch.Tensor, g_shift: torch.Tensor, b_shift: torch.Tensor) -> torch.Tensor:
    """Shift rgb channels.

    Shift each image's channel by either r_shift for red, g_shift for green and b_shift for blue channels.
    """

    KORNIA_CHECK_IS_TENSOR(image)
    KORNIA_CHECK_IS_COLOR(image, f"with shape {image.shape}")

    shifts = [r_shift, g_shift, b_shift]

    shifted = (image + torch.Tensor(shifts).view(1, 3, 1, 1).to(image)).clamp_(min=0, max=1)

    return shifted

kornia/enhance/shift_rgb.py
import torch

from kornia.testing import KORNIA_CHECK_IS_COLOR, KORNIA_CHECK_IS_TENSOR


def shift_rgb(image: torch.Tensor, r_shift: torch.Tensor, g_shift: torch.Tensor, b_shift: torch.Tensor) -> torch.Tensor:
    """Shift rgb channels.

    Shift each image's channel by either r_shift for red, g_shift for green and b_shift for blue channels.
    """

    KORNIA_CHECK_IS_TENSOR(image)
    KORNIA_CHECK_IS_COLOR(image, f"with shape {image.shape}")

    shifts = [r_shift, g_shift, b_shift]

    shifted = (image + torch.stack(shifts).view(-1, 3, 1, 1).to(image)).clamp_(min=0, max=1)

    return shifted

kornia/enhance/shift_rgb.py
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 12cf1c0ba5..faa9fe9e84 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -63,7 +63,7 @@ repos:
           - "B101,B311"
 
   - repo: https://github.com/myint/docformatter
-    rev: v1.5.0-rc1
+    rev: v1.5.0
     hooks:
       - id: docformatter
         args: [--in-place, --wrap-summaries=115, --wrap-descriptions=120]
diff --git a/kornia/enhance/shift_rgb.py b/kornia/enhance/shift_rgb.py
index 3599f96ed0..e68b84385c 100644
--- a/kornia/enhance/shift_rgb.py
+++ b/kornia/enhance/shift_rgb.py
@@ -14,6 +14,6 @@ def shift_rgb(image: torch.Tensor, r_shift: torch.Tensor, g_shift: torch.Tensor,
 
     shifts = [r_shift, g_shift, b_shift]
 
-    shifted = (image + torch.Tensor(shifts).view(1, 3, 1, 1).to(image)).clamp_(min=0, max=1)
+    shifted = (image + torch.stack(shifts).view(-1, 3, 1, 1).to(image)).clamp_(min=0, max=1)
 
     return shifted
diff --git a/test/augmentation/test_augmentation.py b/test/augmentation/test_augmentation.py
index ca23c739dd..194920c4bc 100644
--- a/test/augmentation/test_augmentation.py
+++ b/test/augmentation/test_augmentation.py
@@ -3339,7 +3339,50 @@ def test_module(self, device, dtype):
 
 class TestRandomRGBShift:
     def test_smoke(self, device, dtype):
-        img = torch.rand(1, 3, 4, 5, device=device, dtype=dtype)
+        img = torch.rand(2, 3, 4, 5, device=device, dtype=dtype)
         aug = RandomRGBShift(p=1.0).to(device)
         out = aug(img)
-        assert out.shape == (1, 3, 4, 5)
+        assert out.shape == (2, 3, 4, 5)
+
+    def test_onnx_export(self, device, dtype):
+        img = torch.rand(1, 3, 4, 5, device=device, dtype=dtype)
+        aug = RandomRGBShift(p=1.0).to(device)
+        torch.onnx.export(aug, img, "temp.onnx", export_params=True)
+
+    def test_random_rgb_shift(self, device, dtype):
+        torch.manual_seed(0)
+        input = torch.tensor(
+            [[[[0.2, 0.0]], [[0.3, 0.5]], [[0.4, 0.7]]], [[[0.2, 0.7]], [[0.0, 0.8]], [[0.2, 0.3]]]],
+            device=device,
+            dtype=dtype,
+        )
+
+        f = RandomRGBShift(p=1.0).to(device)
+        expected = torch.tensor(
+            [
+                [[[0.19625, 0.00000]], [[0.56822, 0.76822]], [[0.00000, 0.28847]]],
+                [[[0.00000, 0.33203]], [[0.00000, 0.60742]], [[0.33407, 0.43407]]],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+        utils.assert_close(f(input), expected, rtol=1e-4, atol=1e-4)
+
+    def test_random_rgb_shift_same_batch(self, device, dtype):
+        torch.manual_seed(0)
+        input = torch.tensor(
+            [[[[0.2, 0.0]], [[0.3, 0.5]], [[0.4, 0.7]]], [[[0.2, 0.7]], [[0.0, 0.8]], [[0.2, 0.3]]]],
+            device=device,
+            dtype=dtype,
+        )
+
+        f = RandomRGBShift(p=1.0, same_on_batch=True).to(device)
+        expected = torch.tensor(
+            [
+                [[[0.19626, 0.00000]], [[0.29626, 0.49626]], [[0.66822, 0.96822]]],
+                [[[0.46822, 0.96822]], [[0.00000, 0.38848]], [[0.00000, 0.00000]]],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+        utils.assert_close(f(input), expected, rtol=1e-4, atol=1e-4)
diff --git a/test/enhance/test_shift_rgb.py b/test/enhance/test_shift_rgb.py
index 7a8dcc0d56..484675e034 100644
--- a/test/enhance/test_shift_rgb.py
+++ b/test/enhance/test_shift_rgb.py
@@ -8,7 +8,7 @@
 
 class TestRGBShift:
     def test_rgb_shift_no_shift(self, device, dtype):
-        r_shift, g_shift, b_shift = 0, 0, 0
+        r_shift, g_shift, b_shift = torch.Tensor([0]), torch.Tensor([0]), torch.Tensor([0])
         image = torch.rand(2, 3, 5, 5, device=device, dtype=dtype)
         expected = image
         shifted = kornia.enhance.shift_rgb(image, r_shift, g_shift, b_shift)
@@ -16,7 +16,7 @@ def test_rgb_shift_no_shift(self, device, dtype):
         utils.assert_close(shifted, expected)
 
     def test_rgb_shift_all_zeros(self, device, dtype):
-        r_shift, g_shift, b_shift = -0.1, -0.1, -0.1
+        r_shift, g_shift, b_shift = torch.Tensor([-0.1]), torch.Tensor([-0.1]), torch.Tensor([-0.1])
         image = torch.zeros(2, 3, 5, 5, device=device, dtype=dtype)
         expected = image
         shifted = kornia.enhance.shift_rgb(image, r_shift, g_shift, b_shift)
@@ -24,7 +24,7 @@ def test_rgb_shift_all_zeros(self, device, dtype):
         utils.assert_close(shifted, expected)
 
     def test_rgb_shift_all_ones(self, device, dtype):
-        r_shift, g_shift, b_shift = 1, 1, 1
+        r_shift, g_shift, b_shift = torch.Tensor([1]), torch.Tensor([1]), torch.Tensor([1])
         image = torch.rand(2, 3, 5, 5, device=device, dtype=dtype)
         expected = torch.ones(2, 3, 5, 5, device=device, dtype=dtype)
         shifted = kornia.enhance.shift_rgb(image, r_shift, g_shift, b_shift)
@@ -32,21 +32,29 @@ def test_rgb_shift_all_ones(self, device, dtype):
         utils.assert_close(shifted, expected)
 
     def test_rgb_shift_invalid_parameter_shape(self, device, dtype):
-        r_shift, g_shift, b_shift = 0.5, 0.5, 0.5
+        r_shift, g_shift, b_shift = torch.Tensor([0.5]), torch.Tensor([0.5]), torch.Tensor([0.5])
         image = torch.randn(3, 3, device=device, dtype=dtype)
         with pytest.raises(TypeError):
             kornia.enhance.shift_rgb(image, r_shift, g_shift, b_shift)
 
     def test_rgb_shift_gradcheck(self, device, dtype):
-        r_shift, g_shift, b_shift = 0.4, 0.5, 0.2
+        r_shift, g_shift, b_shift = torch.Tensor([0.4]), torch.Tensor([0.5]), torch.Tensor([0.2])
         image = torch.randn(2, 3, 5, 5, device=device, dtype=dtype)
         image = utils.tensor_to_gradcheck_var(image)  # to var
         assert gradcheck(kornia.enhance.shift_rgb, (image, r_shift, g_shift, b_shift), raise_exception=True)
 
     def test_rgb_shift(self, device, dtype):
-        r_shift, g_shift, b_shift = 0.1, 0.2, -0.3
-        image = torch.tensor([[[[0.2, 0.0]], [[0.3, 0.5]], [[0.4, 0.7]]]], device=device, dtype=dtype)
+        r_shift, g_shift, b_shift = torch.Tensor([0.1]), torch.Tensor([0.3]), torch.Tensor([-0.3])
+        image = torch.tensor(
+            [[[[0.2, 0.0]], [[0.3, 0.5]], [[0.4, 0.7]]], [[[0.2, 0.7]], [[0.0, 0.8]], [[0.2, 0.3]]]],
+            device=device,
+            dtype=dtype,
+        )
         shifted = kornia.enhance.shift_rgb(image, r_shift, g_shift, b_shift)
-        expected = torch.tensor([[[[0.3, 0.1]], [[0.5, 0.7]], [[0.1, 0.4]]]], device=device, dtype=dtype)
+        expected = torch.tensor(
+            [[[[0.3, 0.1]], [[0.6, 0.8]], [[0.1, 0.4]]], [[[0.3, 0.8]], [[0.3, 1.0]], [[0.0, 0.0]]]],
+            device=device,
+            dtype=dtype,
+        )
 
         utils.assert_close(shifted, expected)

Bugs when using RandomRGBShift
### Describe the bug

When running RandomRGBShift augmentation on gpus I came across this error 

```
  File "/usr/local/lib/python3.8/dist-packages/kornia/augmentation/_2d/intensity/random_rgb_shift.py", line 100, in apply_transform
    return shift_rgb(inp, params['r_shift'], params['g_shift'], params['b_shift'])
  File "/usr/local/lib/python3.8/dist-packages/kornia/enhance/shift_rgb.py", line 17, in shift_rgb
    shifted = (image + torch.Tensor(shifts).view(1, 3, 1, 1).to(image)).clamp_(min=0, max=1)
ValueError: only one element tensors can be converted to Python scalars
```
after checking the code I think the problem is that 
```
shifts = [params['r_shift'], params['g_shift'], params['b_shift']]
shifted = (image + torch.Tensor(shifts).view(1, 3, 1, 1).to(image)).clamp_(min=0, max=1)
```
and also I am not sure about `.view(1,3,1,1)` shouldn't it be `.view(-1,3,1,1)`

### Reproduction steps

```bash
1. setting up RandomRGBShift augmentation
2. Using RandomRGBShift augmentation with gpus and batchsize > 1
```


### Expected behavior

```
  File "/usr/local/lib/python3.8/dist-packages/kornia/augmentation/_2d/intensity/random_rgb_shift.py", line 100, in apply_transform
    return shift_rgb(inp, params['r_shift'], params['g_shift'], params['b_shift'])
  File "/usr/local/lib/python3.8/dist-packages/kornia/enhance/shift_rgb.py", line 17, in shift_rgb
    shifted = (image + torch.Tensor(shifts).view(1, 3, 1, 1).to(image)).clamp_(min=0, max=1)
ValueError: only one element tensors can be converted to Python scalars
```

### Environment

```shell
- PyTorch Version : 1.10
- Linux
- How you installed PyTorch (`pip`):
- Python version: 3.8
- cuda:11.3.0-cudnn8
```


### Additional context

_No response_
