Repository: kornia/kornia
Issue category: Bug Report

Issue:
AugmentationSequential explicitly moves the output to the CPU if data_keys is given
### Describe the bug

With the 0.7.4 release, augmentations on the GPU are not possible anymore because the output of the input tensor is always explicitly moved to the CPU.

The problem is that `_detach_tensor_to_cpu` is called explicitly on every tensor in [augment.py](https://github.com/kornia/kornia/blob/main/kornia/augmentation/container/augment.py#L507).

### Reproduction steps
The augmented version of the input gets moved to the CPU:

```python
import torch
import kornia

input = torch.randn(1, 3, 10, 10, device="cuda")
mask = torch.randn(1, 10, 10, device="cuda")
aug_list = kornia.augmentation.AugmentationSequential(kornia.augmentation.RandomHorizontalFlip(p=1))

for a in aug_list(input, mask, data_keys=["input", "mask"]):
    print(a.dtype, a.device)
# torch.float32 cpu
# torch.float32 cuda:0
```

Without `data_keys` and only one input, the output device is as expected:
```python
for a in aug_list(input):
    print(a.dtype, a.device)
# torch.float32 cuda:0
```


### Expected behavior

I would expect the augmented version of the input tensor to reside on the GPU.

### Environment

- PyTorch Version (e.g., 1.0): 2.5.1
- OS (e.g., Linux): Ubuntu
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.12
- CUDA/cuDNN version: 12.4
- GPU models and configuration: NVIDIA 4090 RTX
- Any other relevant information:


### Additional context

Reverting back to 0.7.3 fixes the problem.


Hints:
i can see that this was touched in https://github.com/kornia/kornia/pull/2979 @ashnair1 @shijianjian @johnnv1 do you recall why we need here  `_detach_tensor_to_cpu` instead of something like `_detach_tensor_to_device` ?
It was mainly for passing the tests, due to the randomness handling for CUDA and CPU are different.

I think the bug is because the dict tensors are not handled properly to be move back to CUDA. @ashnair1, can you fix this? We may need to do a patch release after this.

BTW, I think we may need a pre-merge action to run CUDA tests. At least to run it for one-time for each merge. It has happened many times I think that some CUDA errors after all CPU tests passed.
> BTW, I think we may need a pre-merge action to run CUDA tests. At least to run it for one-time for each merge. It has happened many times I think that some CUDA errors after all CPU tests passed.

For that, we need someone to sponsor with access to a machine that has a cuda device.
Will the meta device be helpful? So we mock cuda instead?

https://pytorch.org/docs/stable/meta.html

https://github.com/pytorch/pytorch/issues/61654#issuecomment-879989145
we can try yes
One of Kornia's biggest strengths is fast GPU augmentations. This bug makes 0.7.4 unusable for many including me - are there any updates on a fix?
hey @ashnair1, can you help fix this issue?

Patch:
diff --git a/kornia/augmentation/container/augment.py b/kornia/augmentation/container/augment.py
index a9cad91924..ebd892bf17 100644
--- a/kornia/augmentation/container/augment.py
+++ b/kornia/augmentation/container/augment.py
@@ -507,11 +507,9 @@ def __call__(
                 if output_type == "tensor":
                     self._output_image = _output_image
                     if isinstance(_output_image, dict):
-                        self._output_image[original_keys[idx]] = self._detach_tensor_to_cpu(
-                            _output_image[original_keys[idx]]
-                        )
+                        self._output_image[original_keys[idx]] = _output_image[original_keys[idx]]
                     else:
-                        self._output_image[idx] = self._detach_tensor_to_cpu(_output_image[idx])
+                        self._output_image[idx] = _output_image[idx]
                 elif isinstance(_output_image, dict):
                     self._output_image[original_keys[idx]] = _output_image[original_keys[idx]]
                 else:


Test patch: