Repository: kornia/kornia

Issue:
Fix typing annotation errors
### Describe the bug

Our CI for type checking is showing some errors:
- https://github.com/kornia/kornia/actions/workflows/pr_test_typing.yml
- https://github.com/kornia/kornia/actions/workflows/scheduled_test_typing.yml

### Reproduction steps

```bash
Run the CI or follow in his steps
```

~### Additional context~

~Needs to check if is a problem with our typing annotations, or if it's a full bug of Mypy -- anyway the mypy shouldn't be crashing~

~it is worth checking against mypy's main to see if it has already been fixed as well~


Discussion:
I think it doesn't crash, it just outputs all the errors. And sets an error code to 1, because it found problems

Here is the local run:

```
(pytorch) dmytromishkin@Z000845-845 ~/dev/kornia % make mypy
mypy
kornia/core/check.py:100: error: Variable "typ" is not valid as a type  [valid-type]
        return cast(typ, maybe_obj)
                    ^
kornia/core/check.py:100: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
kornia/augmentation/random_generator/base.py:86: error: Argument 1 to "rsample" of "Distribution" has incompatible type "Tuple[int, ...]"; expected "Size"  [arg-type]
            out = self.dist.rsample(sample_shape)
                                    ^~~~~~~~~~~~
kornia/augmentation/random_generator/base.py:92: error: Argument 1 to "sample" of "Distribution" has incompatible type "Tuple[int, ...]"; expected "Size"  [arg-type]
            out = self.dist.sample(sample_shape)
                                   ^~~~~~~~~~~~
kornia/augmentation/random_generator/base.py:98: error: Argument 1 to "sample_n" of "Distribution" has incompatible type "Union[int, float, Tensor]"; expected "int"  [arg-type]
            out = self.dist.sample_n(n)
                                     ^
kornia/x/trainers.py:78: error: Argument 6 to "__init__" of "Trainer" has incompatible type "CosineAnnealingLR"; expected "_LRScheduler"  [arg-type]
            super().__init__(model, train_dataloader, valid_dataloader, criterion, optimizer, scheduler, config, callbacks)
                                                                                              ^~~~~~~~~
kornia/augmentation/_2d/geometric/base.py:145: error: Incompatible types in assignment (expression has type "Tuple[Any, Any]", variable has type "Optional[List[Any]]")  [assignment]
                size = (size[-2], size[-1])
                       ^~~~~~~~~~~~~~~~~~~~
kornia/augmentation/_2d/geometric/base.py:152: error: Argument "size" to "inverse_transform" of "GeometricAugmentationBase2D" has incompatible type "Optional[List[Any]]"; expected "Optional[Tuple[int, int]]" 
[arg-type]
                output = self.inverse_transform(in_tensor, flags=flags, transform=transform, size=size)
                                                                                                  ^~~~
kornia/augmentation/_2d/geometric/base.py:157: error: Argument "size" to "inverse_transform" of "GeometricAugmentationBase2D" has incompatible type "Optional[List[Any]]"; expected "Optional[Tuple[int, int]]" 
[arg-type]
                    size=size,
                         ^~~~
Found 8 errors in 4 files (checked 341 source files)
make: *** [mypy] Error 1

```
very true, I don't know where I saw that it was crashing 😂 

Resolving patch:
diff --git a/kornia/augmentation/_2d/geometric/base.py b/kornia/augmentation/_2d/geometric/base.py
--- a/kornia/augmentation/_2d/geometric/base.py
+++ b/kornia/augmentation/_2d/geometric/base.py
@@ -138,11 +138,11 @@ def inverse_inputs(
             self._params if params is None else params, flags, **kwargs
         )
 
-        size = None
+        size: Optional[Tuple[int, int]] = None
         if "forward_input_shape" in params:
             # Majorly for cropping functions
-            size = params['forward_input_shape'].tolist()
-            size = (size[-2], size[-1])
+            _size = params['forward_input_shape'].tolist()
+            size = (_size[-2], _size[-1])
 
         # if no augmentation needed
         if not to_apply.any():
diff --git a/kornia/feature/disk/structs.py b/kornia/feature/disk/structs.py
--- a/kornia/feature/disk/structs.py
+++ b/kornia/feature/disk/structs.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any
 
 import torch.nn.functional as F
 
@@ -41,7 +42,7 @@ def y(self) -> Tensor:
         """Accesses the y coordinates of keypoints (along image height)."""
         return self.keypoints[:, 1]
 
-    def to(self, *args, **kwargs) -> DISKFeatures:
+    def to(self, *args: Any, **kwargs: Any) -> DISKFeatures:
         """Calls :func:`torch.Tensor.to` on each tensor to move the keypoints, descriptors and detection scores to
         the specified device and/or data type.
 


Test patch: