Repository: kornia/kornia

Issue:
The RandomPosterize class raises RuntimeError because of conflicting devices.
### Describe the bug

Using the _RandomPosterize_ class on the GPU device raises the following exception:

```
  File "<...>\Python39\lib\site-packages\kornia\enhance\adjust.py", line 509, in _right_shift
    return (input * 255).to(torch.uint8) / (2 ** shift).to(input.dtype) / 255.0
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```

It seems like it's caused by the _bits_ variable not being placed on the same device. It's probably fixed just by placing something like the below code in the _posterize_ function.

```
if bits.device != input.device:
    bits = bits.to(input.device)
```



### Reproduction steps

```bash
import torch
from kornia.augmentation import RandomPosterize
from kornia.enhance import posterize

# On CPU works
transform = RandomPosterize(bits=6, p=1.0)
input = torch.rand(size=(1, 3, 128, 128),  dtype=torch.float32)
out = transform(input)

# On GPU fails with: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
transform = RandomPosterize(bits=6, p=1.0).to('cuda:0')
input = torch.rand(size=(1, 3, 128, 128), device='cuda:0', dtype=torch.float32)
out = transform(input)

# The posterize function fails, because bits and input are on different devices
input = torch.rand(size=(1, 3, 128, 128), device='cuda:0', dtype=torch.float32)
bits = torch.tensor([6], device='cpu', dtype=torch.int32)
out = posterize(input, bits)
```


### Expected behavior

The `RandomPosterize` class should work on all devices.

### Environment
Using Kornia 0.6.2

```shell
PyTorch version: 1.10.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Pro
GCC version: (x86_64-posix-seh, Built by strawberryperl.com project) 8.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.4 (tags/v3.9.4:1f2e308, Apr  6 2021, 13:40:21) [MSC v.1928 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19041-SP0
Is CUDA available: True
CUDA runtime version: 11.5.119
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060
Nvidia driver version: 497.29
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] catatorch==0.1.0
[pip3] info-nce-pytorch==0.1.4
[pip3] numpy==1.20.2
[pip3] pytorch-lightning==1.5.7
[pip3] pytorch-lightning-bolts==0.3.2
[pip3] rmi-pytorch==0.1.1
[pip3] torch==1.10.1+cu113
[pip3] torchaudio==0.10.1+cu113
[pip3] torchinfo==1.5.3
[pip3] torchmetrics==0.6.1
[pip3] torchvision==0.11.2+cu113
[conda] Could not collect
```


### Additional context

_No response_


Discussion:
Hi! Do you have the full log error?
Full output for the reproduction script:

```
/src/posterize.py
Traceback (most recent call last):
  File "<...>\src\posterize.py", line 13, in <module>
    out = transform(input)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\augmentation\base.py", line 270, in forward
    output = self.apply_func(in_tensor, in_transform, self._params, return_transform)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\augmentation\base.py", line 229, in apply_func
    output = self.apply_transform(in_tensor, params, trans_matrix)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\augmentation\augmentation.py", line 1543, in apply_transform
    return posterize(input, params["bits_factor"])
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\utils\image.py", line 227, in _wrapper
    output: torch.Tensor = f(input, *args, **kwargs)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\enhance\adjust.py", line 521, in posterize
    return _posterize_one(input, bits)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\enhance\adjust.py", line 518, in _posterize_one
    return _left_shift(_right_shift(input, bits), bits)
  File "C:\Users\Robin\AppData\Local\Programs\Python\Python39\lib\site-packages\kornia\enhance\adjust.py", line 509, in _right_shift
    return (input * 255).to(torch.uint8) / (2 ** shift).to(input.dtype) / 255.0
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```
Awesome, I’ll fix during the weekend 

Resolving patch:
diff --git a/kornia/augmentation/_2d/intensity/posterize.py b/kornia/augmentation/_2d/intensity/posterize.py
--- a/kornia/augmentation/_2d/intensity/posterize.py
+++ b/kornia/augmentation/_2d/intensity/posterize.py
@@ -57,9 +57,10 @@ def __init__(
         keepdim: bool = False,
     ) -> None:
         super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
+        # TODO: the generator should receive the device
         self._param_generator = cast(rg.PosterizeGenerator, rg.PosterizeGenerator(bits))
 
     def apply_transform(
         self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
-        return posterize(input, params["bits_factor"])
+        return posterize(input, params["bits_factor"].to(input.device))
diff --git a/kornia/augmentation/container/utils.py b/kornia/augmentation/container/utils.py
--- a/kornia/augmentation/container/utils.py
+++ b/kornia/augmentation/container/utils.py
@@ -292,6 +292,7 @@ def pad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
         """
         if len(input.shape) not in (3, 4,):
             raise AssertionError(input.shape)
+
         if len(padding_size.shape) != 2:
             raise AssertionError(padding_size.shape)
 
@@ -318,6 +319,7 @@ def unpad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
         """
         if len(input.shape) not in (3, 4,):
             raise AssertionError(input.shape)
+
         if len(padding_size.shape) != 2:
             raise AssertionError(padding_size.shape)
 
@@ -354,7 +356,7 @@ def apply_trans(
 
         padding_size = cls._get_padding_size(module, param)
         if padding_size is not None:
-            _input = cls.pad(_input, padding_size)
+            _input = cls.pad(_input, padding_size.to(_input))
 
         _input, label = super().apply_trans(_input, label, module, param)
 
@@ -378,7 +380,7 @@ def inverse(cls, input: torch.Tensor, module: nn.Module, param: Optional[ParamIt
 
         padding_size = cls._get_padding_size(module, param)
         if padding_size is not None:
-            _input = cls.unpad(_input, padding_size)
+            _input = cls.unpad(_input, padding_size.to(input))
 
         return _input
 
@@ -393,16 +395,18 @@ class BBoxXYXYApplyInverse(BBoxApplyInverse):
 
     @classmethod
     def pad(cls, input, padding_size):
-        for i in range(len(padding_size)):
-            input[i, :, 0::2] += padding_size[i][0]  # left padding
-            input[i, :, 1::2] += padding_size[i][2]  # top padding
+        _padding_size = padding_size.to(input)
+        for i in range(len(_padding_size)):
+            input[i, :, 0::2] += _padding_size[i][0]  # left padding
+            input[i, :, 1::2] += _padding_size[i][2]  # top padding
         return input
 
     @classmethod
     def unpad(cls, input, padding_size):
-        for i in range(len(padding_size)):
-            input[i, :, 0::2] -= padding_size[i][0]  # left padding
-            input[i, :, 1::2] -= padding_size[i][2]  # top padding
+        _padding_size = padding_size.to(input)
+        for i in range(len(_padding_size)):
+            input[i, :, 0::2] -= _padding_size[i][0]  # left padding
+            input[i, :, 1::2] -= _padding_size[i][2]  # top padding
         return input
 
     @classmethod
@@ -428,18 +432,20 @@ class BBoxXYWHApplyInverse(BBoxXYXYApplyInverse):
 
     @classmethod
     def pad(cls, input, padding_size):
+        _padding_size = padding_size.to(input)
         # pad only xy, not wh
-        for i in range(len(padding_size)):
-            input[i, :, 0] += padding_size[i][0]  # left padding
-            input[i, :, 1] += padding_size[i][2]  # top padding
+        for i in range(len(_padding_size)):
+            input[i, :, 0] += _padding_size[i][0]  # left padding
+            input[i, :, 1] += _padding_size[i][2]  # top padding
         return input
 
     @classmethod
     def unpad(cls, input, padding_size):
+        _padding_size = padding_size.to(input)
         # unpad only xy, not wh
-        for i in range(len(padding_size)):
-            input[i, :, 0] -= padding_size[i][0]  # left padding
-            input[i, :, 1] -= padding_size[i][2]  # top padding
+        for i in range(len(_padding_size)):
+            input[i, :, 0] -= _padding_size[i][0]  # left padding
+            input[i, :, 1] -= _padding_size[i][2]  # top padding
         return input
 
 
@@ -457,6 +463,7 @@ def pad(cls, input: torch.Tensor, padding_size: torch.Tensor) -> torch.Tensor:
 
         if len(input.shape) not in (2, 3,):
             raise AssertionError(input.shape)
+
         if len(padding_size.shape) != 2:
             raise AssertionError(padding_size.shape)
 
diff --git a/kornia/augmentation/random_generator/_2d/posterize.py b/kornia/augmentation/random_generator/_2d/posterize.py
--- a/kornia/augmentation/random_generator/_2d/posterize.py
+++ b/kornia/augmentation/random_generator/_2d/posterize.py
@@ -47,7 +47,7 @@ def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None:
     def forward(self, batch_shape: torch.Size, same_on_batch: bool = False) -> Dict[str, torch.Tensor]:  # type:ignore
         batch_size = batch_shape[0]
         _common_param_check(batch_size, same_on_batch)
-        _device, _dtype = _extract_device_dtype([self.bits if isinstance(self.bits, torch.Tensor) else None])
+        _device, _ = _extract_device_dtype([self.bits if isinstance(self.bits, torch.Tensor) else None])
         bits_factor = _adapted_rsampling((batch_size,), self.bit_sampler, same_on_batch)
         return dict(bits_factor=bits_factor.to(device=_device, dtype=torch.int32))
 


Test patch:
diff --git a/test/augmentation/test_augmentation.py b/test/augmentation/test_augmentation.py
--- a/test/augmentation/test_augmentation.py
+++ b/test/augmentation/test_augmentation.py
@@ -25,6 +25,7 @@
     RandomGrayscale,
     RandomHorizontalFlip,
     RandomInvert,
+    RandomPosterize,
     RandomResizedCrop,
     RandomRotation,
     RandomThinPlateSpline,
@@ -2780,3 +2781,11 @@ def test_smoke(self, device, dtype):
         out = aug(img)
         assert out.shape == (1, 1, 4, 5)
         assert (aug.inverse(out) == img).all()
+
+
+class TestRandomPosterize:
+    def test_smoke(self, device, dtype):
+        img = torch.rand(1, 1, 4, 5, device=device, dtype=dtype)
+        aug = RandomPosterize(bits=6, p=1.0).to(device)
+        out = aug(img)
+        assert out.shape == (1, 1, 4, 5)
diff --git a/test/augmentation/test_container.py b/test/augmentation/test_container.py
--- a/test/augmentation/test_container.py
+++ b/test/augmentation/test_container.py
@@ -520,6 +520,8 @@ def test_forward_and_inverse_return_transform(self, random_apply, device, dtype)
 
         reproducibility_test((inp, mask, bbox, keypoints), aug)
 
+        # TODO(jian): we sometimes throw the following error
+        # AttributeError: 'tuple' object has no attribute 'shape'
         out_inv = aug.inverse(*out)
         assert out_inv[0].shape == inp.shape
         assert out_inv[1].shape == mask.shape