Repository: kornia/kornia

Issue:
Batched Augmentations Fail With `torch.autocast()` When `0 < p < 1`
### Describe the bug

When running batched augmentations in mixed precision, kornia bugs out with the following error:

```python
RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and Half for the source.
```

Interestingly, this bug doesn't happen when `p=0.0` or `p=1.0`, but only when it's between these. My hunch is there's some erroneous dangling hidden state that's causing this.


### Reproduction steps

```python
import torch
import kornia as K

aug = K.augmentation.RandomGaussianBlur((3,3), (0.1, 3), p=0.0)
x = torch.rand(2,3,384,640).cuda()

for i in range(100):
    with torch.autocast("cuda", dtype=torch.half):
        aug(x)
print("Passed with p=0.0")


aug = K.augmentation.RandomGaussianBlur((3,3), (0.1, 3), p=1.0)
for i in range(100):
    with torch.autocast("cuda", dtype=torch.half):
        aug(x)
print("Passed with p=1.0")


aug = K.augmentation.RandomGaussianBlur((3,3), (0.1, 3), p=0.5)
for i in range(100):
    with torch.autocast("cuda", dtype=torch.half):
        aug(x)
print("Passed with p=0.5")
```


### Expected behavior

Run augmentations without any errors

### Environment

Inside a Docker environment, built from the base NVIDIA NGC image `nvcr.io/nvidia/pytorch:21.12-py3`

```shell
PyTorch version: 1.11.0a0+b6df043
Is debug build: False
CUDA used to build PyTorch: 11.5
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.21.3
Libc version: glibc-2.31

Python version: 3.8.12 | packaged by conda-forge | (default, Oct 12 2021, 21:59:51)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-44-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.5.50
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090

Nvidia driver version: 510.73.08
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] nvidia-dlprof-pytorch-nvtx==1.8.0
[pip3] pytorch-lightning==1.6.0
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.11.0a0+b6df043
[pip3] torch-tensorrt==1.1.0a0
[pip3] torchmetrics==0.7.0
[pip3] torchtext==0.12.0a0
[pip3] torchvision==0.11.0a0
[conda] magma-cuda110             2.5.2                         5    local
[conda] mkl                       2019.5                      281    conda-forge
[conda] mkl-include               2019.5                      281    conda-forge
[conda] numpy                     1.21.2                   pypi_0    pypi
[conda] nvidia-dlprof-pytorch-nvtx 1.8.0                    pypi_0    pypi
[conda] pytorch-lightning         1.6.0                    pypi_0    pypi
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.11.0a0+b6df043          pypi_0    pypi
[conda] torch-tensorrt            1.1.0a0                  pypi_0    pypi
[conda] torchmetrics              0.7.0                    pypi_0    pypi
[conda] torchtext                 0.12.0a0                 pypi_0    pypi
[conda] torchvision               0.11.0a0                 pypi_0    pypi
```


### Additional context

This issue was originally opened in pytorch-lightning (https://github.com/PyTorchLightning/pytorch-lightning/issues/13228) , but as they pointed out, the bug is in kornia. 

This _may_ potentially be the same issue as https://github.com/kornia/kornia/issues/1477, which was closed thinking the bug was in lightning. 


Batched Augmentations Fail With `torch.autocast()` When `0 < p < 1`
### Describe the bug

When running batched augmentations in mixed precision, kornia bugs out with the following error:

```python
RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and Half for the source.
```

Interestingly, this bug doesn't happen when `p=0.0` or `p=1.0`, but only when it's between these. My hunch is there's some erroneous dangling hidden state that's causing this.


### Reproduction steps

```python
import torch
import kornia as K

aug = K.augmentation.RandomGaussianBlur((3,3), (0.1, 3), p=0.0)
x = torch.rand(2,3,384,640).cuda()

for i in range(100):
    with torch.autocast("cuda", dtype=torch.half):
        aug(x)
print("Passed with p=0.0")


aug = K.augmentation.RandomGaussianBlur((3,3), (0.1, 3), p=1.0)
for i in range(100):
    with torch.autocast("cuda", dtype=torch.half):
        aug(x)
print("Passed with p=1.0")


aug = K.augmentation.RandomGaussianBlur((3,3), (0.1, 3), p=0.5)
for i in range(100):
    with torch.autocast("cuda", dtype=torch.half):
        aug(x)
print("Passed with p=0.5")
```


### Expected behavior

Run augmentations without any errors

### Environment

Inside a Docker environment, built from the base NVIDIA NGC image `nvcr.io/nvidia/pytorch:21.12-py3`

```shell
PyTorch version: 1.11.0a0+b6df043
Is debug build: False
CUDA used to build PyTorch: 11.5
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.21.3
Libc version: glibc-2.31

Python version: 3.8.12 | packaged by conda-forge | (default, Oct 12 2021, 21:59:51)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.13.0-44-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.5.50
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090

Nvidia driver version: 510.73.08
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] nvidia-dlprof-pytorch-nvtx==1.8.0
[pip3] pytorch-lightning==1.6.0
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.11.0a0+b6df043
[pip3] torch-tensorrt==1.1.0a0
[pip3] torchmetrics==0.7.0
[pip3] torchtext==0.12.0a0
[pip3] torchvision==0.11.0a0
[conda] magma-cuda110             2.5.2                         5    local
[conda] mkl                       2019.5                      281    conda-forge
[conda] mkl-include               2019.5                      281    conda-forge
[conda] numpy                     1.21.2                   pypi_0    pypi
[conda] nvidia-dlprof-pytorch-nvtx 1.8.0                    pypi_0    pypi
[conda] pytorch-lightning         1.6.0                    pypi_0    pypi
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.11.0a0+b6df043          pypi_0    pypi
[conda] torch-tensorrt            1.1.0a0                  pypi_0    pypi
[conda] torchmetrics              0.7.0                    pypi_0    pypi
[conda] torchtext                 0.12.0a0                 pypi_0    pypi
[conda] torchvision               0.11.0a0                 pypi_0    pypi
```


### Additional context

This issue was originally opened in pytorch-lightning (https://github.com/PyTorchLightning/pytorch-lightning/issues/13228) , but as they pointed out, the bug is in kornia. 

This _may_ potentially be the same issue as https://github.com/kornia/kornia/issues/1477, which was closed thinking the bug was in lightning. 




Discussion:
We might need to dig into amp to understand how pytorch cast/uncast types. /cc @shijianjian 
@rsomani95 feel free to propose a solution— until we don’t get any kind of cuda sponsorship I don’t think we’ll prioritise this feature. I can take a look myself at this feature when I need it to optimise some internal projects possibly after summer, but for now I don’t require half precision to train our models.
@edgarriba Fair enough, thanks for the quick response. I'd love to be able to help but I'm out of my depth with how pytorch is casting types. If I make any progress / workaround, I'll definitely update on this thread.
We might need to dig into amp to understand how pytorch cast/uncast types. /cc @shijianjian 
@rsomani95 feel free to propose a solution— until we don’t get any kind of cuda sponsorship I don’t think we’ll prioritise this feature. I can take a look myself at this feature when I need it to optimise some internal projects possibly after summer, but for now I don’t require half precision to train our models.
@edgarriba Fair enough, thanks for the quick response. I'd love to be able to help but I'm out of my depth with how pytorch is casting types. If I make any progress / workaround, I'll definitely update on this thread.

Resolving patch:
diff --git a/kornia/augmentation/_2d/mix/base.py b/kornia/augmentation/_2d/mix/base.py
--- a/kornia/augmentation/_2d/mix/base.py
+++ b/kornia/augmentation/_2d/mix/base.py
@@ -63,7 +63,7 @@ def transform_input(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[
             output = self.apply_non_transform(in_tensor, params, flags)
             output = output.index_put((to_apply,), self.apply_non_transform(applied, params, flags))
         output = _transform_output_shape(output, ori_shape) if self.keepdim else output
-        return output
+        return output.type(in_tensor.dtype)
 
     def transform_mask(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor:
         to_apply = params['batch_prob']
diff --git a/kornia/augmentation/base.py b/kornia/augmentation/base.py
--- a/kornia/augmentation/base.py
+++ b/kornia/augmentation/base.py
@@ -246,6 +246,14 @@ def apply_func(
             flags = self.flags
         to_apply = params['batch_prob']
 
+        # Note about the explicit type conversions here:
+        # If this code is run in an autocast-enabled region, the transformation matrix or the output
+        # may be float16 even though the input was float32 (e.g. torch.mm produces float16 output),
+        # see also the documentation on autocasting: https://pytorch.org/docs/stable/amp.html.
+        # It may be unexpected for the user if the output type changes and it can also lead to errors
+        # for the index_put operation below (see also https://github.com/kornia/kornia/issues/1737)
+        # In case the type already matches the input type, the conversions are no-ops
+
         # if no augmentation needed
         if not to_apply.any():
             output = in_tensor
@@ -258,13 +266,16 @@ def apply_func(
             output = in_tensor.clone()
             trans_matrix = self.identity_matrix(in_tensor)
             trans_matrix = trans_matrix.index_put(
-                (to_apply,), self.compute_transformation(in_tensor[to_apply], params=params, flags=flags)
+                (to_apply,),
+                self.compute_transformation(in_tensor[to_apply], params=params, flags=flags).type(in_tensor.dtype),
             )
             output = output.index_put(
                 (to_apply,),
-                self.apply_transform(in_tensor[to_apply], params=params, flags=flags, transform=trans_matrix[to_apply]),
+                self.apply_transform(
+                    in_tensor[to_apply], params=params, flags=flags, transform=trans_matrix[to_apply]
+                ).type(in_tensor.dtype),
             )
 
-        self._transform_matrix = trans_matrix
+        self._transform_matrix = trans_matrix.type(in_tensor.dtype)
 
-        return output
+        return output.type(in_tensor.dtype)


Test patch:
diff --git a/test/augmentation/test_base.py b/test/augmentation/test_base.py
--- a/test/augmentation/test_base.py
+++ b/test/augmentation/test_base.py
@@ -5,6 +5,7 @@
 from torch.autograd import gradcheck
 
 import kornia.testing as utils  # test utils
+from kornia.augmentation import RandomGaussianBlur
 from kornia.augmentation._2d.base import AugmentationBase2D
 from kornia.augmentation.base import _BasicAugmentationBase
 from kornia.testing import assert_close
@@ -75,6 +76,28 @@ def test_forward(self, device, dtype, keepdim):
             assert output.shape == expected_output.shape
             assert_close(output, expected_output)
 
+    @pytest.mark.parametrize("batch_prob", [[True, True], [False, True], [False, False]])
+    def test_autocast(self, batch_prob, device, dtype):
+        if not hasattr(torch, "autocast"):
+            pytest.skip("PyTorch version without autocast support")
+
+        torch.manual_seed(42)
+
+        aug = RandomGaussianBlur((3, 3), (0.1, 3), p=1)
+        x = torch.rand(len(batch_prob), 3, 100, 100, dtype=dtype).to(device)
+
+        # Explicitly set the batch probability to trigger the different conditions in apply_func()
+        params = aug.forward_parameters(x.shape)
+        params["batch_prob"] = torch.tensor(batch_prob, device=params["batch_prob"].device)
+
+        # There is usually only a sigma value for the batch images which get transformed
+        params["sigma"] = params["sigma"][params["batch_prob"]]
+
+        with torch.autocast(device.type):
+            res = aug(x, params)
+
+        assert res.dtype == dtype, "The output dtype should match the input dtype"
+
 
 class TestAugmentationBase2D:
     @pytest.mark.parametrize(