Repository: kornia/kornia

Issue:
Differentiable JPEG function does not work with GPU tensors
### Describe the bug

Got different devices (CPU and GPU) error while using `kornia.enhance.jpeg_codec_differentiable` with GPU input tensor

### Reproduction steps

```bash
1.

pip install git+https://github.com/kornia/kornia

2.

a = torch.rand(4, 3, 112, 112).to('cuda:0')
kornia.enhance.jpeg_codec_differentiable(
    input=a, 
    jpeg_quality=torch.full([a.shape[0]], 15).to(a.device))

3.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[22], line 2
      1 a = torch.rand(4, 3, 112, 112).to('cuda:0')
----> 2 kornia.enhance.jpeg_codec_differentiable(
      3     input=a, 
      4     jpeg_quality=torch.full([a.shape[0]], 15).to(a.device))

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/utils/image.py:231, in perform_keep_shape_image.<locals>._wrapper(input, *args, **kwargs)
    229 input_shape = input.shape
    230 input = _to_bchw(input)  # view input as (B, C, H, W)
--> 231 output = f(input, *args, **kwargs)
    232 if len(input_shape) == 3:
    233     output = output[0]

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/enhance/jpeg.py:484, in jpeg_codec_differentiable(image_rgb, jpeg_quality, quantization_table_y, quantization_table_c)
    482 quantization_table_c = quantization_table_c.to(device, dtype)
    483 # Perform encoding
--> 484 y_encoded, cb_encoded, cr_encoded = _jpeg_encode(
    485     image_rgb=image_rgb,
    486     jpeg_quality=jpeg_quality,
    487     quantization_table_c=quantization_table_c,
    488     quantization_table_y=quantization_table_y,
    489 )
    490 image_rgb_jpeg: Tensor = _jpeg_decode(
    491     input_y=y_encoded,
    492     input_cb=cb_encoded,
   (...)
    498     quantization_table_y=quantization_table_y,
    499 )
    500 # Clip coded image

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/enhance/jpeg.py:276, in _jpeg_encode(image_rgb, jpeg_quality, quantization_table_y, quantization_table_c)
    274 image_ycbcr = 255.0 * image_ycbcr
    275 # Perform chroma subsampling
--> 276 input_y, input_cb, input_cr = _chroma_subsampling(image_ycbcr)
    277 # Patchify, DCT, and rounding
    278 input_y, input_cb, input_cr = _patchify_8x8(input_y), _patchify_8x8(input_cb), _patchify_8x8(input_cr)

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/enhance/jpeg.py:231, in _chroma_subsampling(input_ycbcr)
    229 output_cr: Tensor = input_ycbcr[:, 2]
    230 # Perform average pooling of Cb and Cr channels
--> 231 output_cb = rescale(output_cb[:, None], factor=0.5, interpolation="bilinear", align_corners=False, antialias=True)
    232 output_cr = rescale(output_cr[:, None], factor=0.5, interpolation="bilinear", align_corners=False, antialias=True)
    233 return output_y, output_cb[:, 0], output_cr[:, 0]

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/geometry/transform/affwarp.py:643, in rescale(input, factor, interpolation, align_corners, antialias)
    641 height, width = input.size()[-2:]
    642 size = (int(height * factor_vert), int(width * factor_horz))
--> 643 return resize(input, size, interpolation=interpolation, align_corners=align_corners, antialias=antialias)

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/utils/image.py:231, in perform_keep_shape_image.<locals>._wrapper(input, *args, **kwargs)
    229 input_shape = input.shape
    230 input = _to_bchw(input)  # view input as (B, C, H, W)
--> 231 output = f(input, *args, **kwargs)
    232 if len(input_shape) == 3:
    233     output = output[0]

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/geometry/transform/affwarp.py:598, in resize(input, size, interpolation, align_corners, side, antialias)
    595     if (ks[1] % 2) == 0:
    596         ks = ks[0], ks[1] + 1
--> 598     input = gaussian_blur2d(input, ks, sigmas)
    600 output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
    601 return output

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/filters/gaussian.py:63, in gaussian_blur2d(input, kernel_size, sigma, border_type, separable)
     61 ky, kx = _unpack_2d_ks(kernel_size)
     62 bs = sigma.shape[0]
---> 63 kernel_x = get_gaussian_kernel1d(kx, sigma[:, 1].view(bs, 1))
     64 kernel_y = get_gaussian_kernel1d(ky, sigma[:, 0].view(bs, 1))
     65 out = filter2d_separable(input, kernel_x, kernel_y, border_type)

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/filters/kernels.py:535, in get_gaussian_kernel1d(kernel_size, sigma, force_even, device, dtype)
    512 r"""Function that returns Gaussian filter coefficients.
    513 
    514 Args:
   (...)
    531             [0.0096, 0.2054, 0.5699, 0.2054, 0.0096]])
    532 """
    533 _check_kernel_size(kernel_size, allow_even=force_even)
--> 535 return gaussian(kernel_size, sigma, device=device, dtype=dtype)

File ~/miniconda3/envs/pytorch310/lib/python3.10/site-packages/kornia/filters/kernels.py:97, in gaussian(window_size, sigma, mean, device, dtype)
     94 KORNIA_CHECK_IS_TENSOR(mean)
     95 KORNIA_CHECK_SHAPE(mean, ["B", "1"])
---> 97 x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - mean).expand(batch_size, -1)
     99 if window_size % 2 == 0:
    100     x = x + 0.5

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```


### Expected behavior

Expected to have no errors while passing two tensors on the same device

### Environment

```shell
It doesn't matter in this case
```


### Additional context

I open this issue for creating a pull request with fix


Discussion:

Resolving patch:
diff --git a/kornia/filters/kernels.py b/kornia/filters/kernels.py
--- a/kornia/filters/kernels.py
+++ b/kornia/filters/kernels.py
@@ -89,7 +89,7 @@ def gaussian(
 
     mean = float(window_size // 2) if mean is None else mean
     if isinstance(mean, float):
-        mean = tensor([[mean]], device=device, dtype=dtype)
+        mean = tensor([[mean]], device=sigma.device, dtype=sigma.dtype)
 
     KORNIA_CHECK_IS_TENSOR(mean)
     KORNIA_CHECK_SHAPE(mean, ["B", "1"])


Test patch: