-
-
Notifications
You must be signed in to change notification settings - Fork 942
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
RandomJPEG
can't handle input on different from CPU device.
#2867
Comments
the augmentations in the end are nn.Module 's, this same behaviour i believe you face when you forward a tensor to a regular model in pytorch that the tensor should match with the params device and not the other way around @johnnv1 @shijianjian |
Yeah, It should work like with import torch
from kornia.augmentation import RandomJPEG
device = "cuda"
jpegq = (1.0, 50.0)
aug = RandomJPEG(jpeg_quality=jpegq, p=1.0).to(device)
example_input = torch.randn((3, 224, 224)).to(device)
res = aug(example_input) And here is an error about the wrong devices: /home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3483.)
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
Traceback (most recent call last):
File "/home/dmdr/Documents/Code/Python/aaa/ptrainer/tmp.py", line 27, in <module>
res = aug(example_input)
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/augmentation/base.py", line 210, in forward
output = self.apply_func(in_tensor, params, flags)
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/augmentation/_2d/base.py", line 129, in apply_func
output = self.transform_inputs(in_tensor, params, flags, trans_matrix)
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/augmentation/base.py", line 261, in transform_inputs
output = self.apply_transform(in_tensor, params, flags, transform=transform)
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/augmentation/_2d/intensity/jpeg.py", line 56, in apply_transform
jpeg_output: Tensor = jpeg_codec_differentiable(input, params["jpeg_quality"])
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/utils/image.py", line 231, in _wrapper
output = f(input, *args, **kwargs)
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/enhance/jpeg.py", line 484, in jpeg_codec_differentiable
y_encoded, cb_encoded, cr_encoded = _jpeg_encode(
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/enhance/jpeg.py", line 281, in _jpeg_encode
y_encoded: Tensor = _quantize(
File "/home/dmdr/miniconda3/envs/ptrain/lib/python3.10/site-packages/kornia/enhance/jpeg.py", line 177, in _quantize
quantization_table[:, None] * _jpeg_quality_to_scale(jpeg_quality)[:, None, None, None]
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! |
@ditwoo thanks ! we'll try to fix unless you want to give it a shot |
@edgarriba I can write a PR with a fix. |
@ditwoo thanks ! very appreciated |
Describe the bug
RandomJPEG
throws an error when the input tensor is not on the CPU. To somehow handle this issue need to manually pass a tensor with the right device to thejpeg_quality
parameter whenRandomJPEG
is created.Reproduction steps
An error:
Expected behavior
I expect that
RandomJPEG
will understand the tensor location (device) and do all the operations on a device that was passed as an argument likeRandomAffine
or any other augmentation does.Environment
Additional context
No response
The text was updated successfully, but these errors were encountered: