Repository: kornia/kornia

Issue:
`NotImplementedError` for elastic transformation with probability p < 1
### Describe the bug

With the newest kornia release (0.6.11), the random elastic transformation fails if it is not applied to every image in the batch.

The problem is that the `apply_non_transform_mask()` method in `_AugmentationBase` per default raises an `NotImplementedError` and since this method is not overwritten in `RandomElasticTransform`, the error is raised. I see that for the other `apply_non*` methods the default is to just return the input.

I see two different solutions:
1. Change the default for `apply_non_transform_mask` to return the input in `_AugmentationBase`.
2. Overwrite the method in `RandomElasticTransform` and just return the input there.

There might be good reasons to keep the `NotImplementedError` in the base class, therefore I wanted to ask first what solution you prefer. I could make a PR for this.

### Reproduction steps

```python
import torch
import kornia.augmentation as K

features = torch.rand(5, 100, 480, 640, dtype=torch.float32, device="cuda")
labels = torch.randint(0, 10, (5, 1, 480, 640), dtype=torch.int64, device="cuda")

torch.manual_seed(0)
aug = K.AugmentationSequential(
    K.RandomElasticTransform(alpha=(0.7, 0.7), sigma=(16, 16), padding_mode="reflection", p=0.2)
)

features_transformed, labels_transformed = aug(features, labels.float(), data_keys=["input", "mask"])
```


### Expected behavior

No `NotImplementedError`.

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
- PyTorch Version (e.g., 1.0): 2.0
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.10.9
- CUDA/cuDNN version: 11.8
- GPU models and configuration: 3090
- Any other relevant information:
```


Discussion:
@JanSellner Please overwrite the `apply_non_transform_mask` of `RandomElasticTransform`.

Resolving patch:
diff --git a/kornia/augmentation/_2d/geometric/elastic_transform.py b/kornia/augmentation/_2d/geometric/elastic_transform.py
--- a/kornia/augmentation/_2d/geometric/elastic_transform.py
+++ b/kornia/augmentation/_2d/geometric/elastic_transform.py
@@ -22,8 +22,6 @@ class RandomElasticTransform(AugmentationBase2D):
           in the y and x directions, respectively.
         align_corners: Interpolation flag used by `grid_sample`.
         resample: Interpolation mode used by `grid_sample`. Either 'nearest' (0) or 'bilinear' (1).
-        mode: Deprecated: Interpolation mode used by `grid_sample`. Either 'bilinear' or 'nearest'.
-          Please use the `resample` argument instead.
         padding_mode: The padding used by ```grid_sample```. Either 'zeros', 'border' or 'refection'.
         same_on_batch: apply the same transformation across the batch.
         p: probability of applying the transformation.
@@ -92,6 +90,11 @@ def apply_transform(
             flags["padding_mode"],
         )
 
+    def apply_non_transform_mask(
+        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
+    ) -> Tensor:
+        return input
+
     def apply_transform_mask(
         self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
     ) -> Tensor:


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
@@ -1,4 +1,5 @@
 from typing import Any, Dict, Optional, Tuple, Type
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -3704,6 +3705,29 @@ def test_mask_transform(self, device, dtype):
         labels_transformed = compose(features, labels, data_keys=["input", "mask"])[1]
         assert_close(labels_transformed.unique(), torch.tensor([0, 10], dtype=dtype, device=device))
 
+    @pytest.mark.parametrize("batch_prob", [[True, True], [False, True], [False, False]])
+    def test_apply(self, batch_prob, device, dtype):
+        torch.manual_seed(0)
+
+        aug_list = AugmentationSequential(RandomElasticTransform(sigma=(2, 2), alpha=(2, 2)))
+        features = torch.rand(2, 3, 10, 10, dtype=dtype, device=device)
+        labels = torch.randint(0, 10, (2, 1, 10, 10), dtype=dtype, device=device)
+
+        # Make sure the transformation works correctly even if only applied to some images in the batch
+        to_apply = torch.tensor(batch_prob, device=device)
+        with patch.object(aug_list[0], '__batch_prob_generator__', return_value=to_apply):
+            features_transformed, labels_transformed = aug_list(features, labels, data_keys=["input", "mask"])
+            assert_close(aug_list._params[0].data["batch_prob"], to_apply)
+
+            # Images should remain unchanged if the transformation is not applied
+            assert_close(features_transformed[~to_apply], features[~to_apply])
+            assert_close(labels_transformed[~to_apply], labels[~to_apply])
+
+            # At least one value in the images should change if the transformation is applied
+            if to_apply.any():
+                assert features_transformed[to_apply].ne(features[to_apply]).any()
+                assert labels_transformed[to_apply].ne(labels[to_apply]).any()
+
 
 class TestRandomThinPlateSpline:
     def test_smoke(self, device, dtype):