Repository: kornia/kornia

Issue:
Loftr does not work with some image size (not a memory issue)
### Describe the bug

LoFTR incorrectly does something with positional embeddings
```
RuntimeError                              Traceback (most recent call last)
[<ipython-input-1-54d246337ab1>](https://9t3p2yszpxn-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20220613-060046-RC00_454553376#) in <module>()
     10                 "image1": torch.rand(1,1, 1704, 2272).cuda()}
     11 with torch.no_grad():
---> 12     correspondences = matcher(input_dict)

3 frames
[/usr/local/lib/python3.7/dist-packages/kornia/feature/loftr/utils/position_encoding.py](https://9t3p2yszpxn-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20220613-060046-RC00_454553376#) in forward(self, x)
     39             x: [N, C, H, W]
     40         """
---> 41         return x + self.pe[:, :, :x.size(2), :x.size(3)]

RuntimeError: The size of tensor a (284) must match the size of tensor b (256) at non-singleton dimension 3
```

### Reproduction steps

```bash
import kornia as K
import kornia.feature as KF
import numpy as np
import torch

matcher = KF.LoFTR(pretrained='outdoor').cuda()

input_dict = {"image0": torch.rand(1,1, 1704, 2272),
                "image1": torch.rand(1,1, 1704, 2272)}
with torch.no_grad():
    correspondences = matcher(input_dict)
```


### Expected behavior

Not an error 

### Environment

```shell
not relevant
```


### Additional context

_No response_


Discussion:
I think this is because the LoFTR positional encoding has a max image size of 2048 pixels by default in a precomputed tensor for the encoding. I guess there are several ways to solve this, one is to let `PositionEncodingSine` recompute a larger positional encoding, and possibly "cache" it, when a larger image size is input.
@georg-bn are you willing to help with such improvement ? there's an ongoing effort to make also loftr scriptable https://github.com/kornia/kornia/pull/1758
Yes, I can have a look.

Resolving patch:
diff --git a/kornia/feature/loftr/utils/position_encoding.py b/kornia/feature/loftr/utils/position_encoding.py
--- a/kornia/feature/loftr/utils/position_encoding.py
+++ b/kornia/feature/loftr/utils/position_encoding.py
@@ -17,25 +17,51 @@ def __init__(self, d_model, max_shape=(256, 256), temp_bug_fix=True):
                 We will remove the buggy impl after re-training all variants of our released models.
         """
         super().__init__()
+        self.d_model = d_model
+        self.temp_bug_fix = temp_bug_fix
 
-        pe = torch.zeros((d_model, *max_shape))
+        pe = self._create_position_encoding(max_shape)
+        self.register_buffer('pe', pe, persistent=False)  # [1, C, H, W]
+
+    def _create_position_encoding(self, max_shape):
+        """Creates a position encoding from scratch.
+
+        For 1/8 feature map (which is standard): If the input image size is H, W (both divisible by 8), the max_shape
+        should be (H//8, W//8).
+        """
+        pe = torch.zeros((self.d_model, *max_shape))
         y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0)
         x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0)
-        if temp_bug_fix:
-            div_term = torch.exp(torch.arange(0, d_model // 2, 2).float() * (-math.log(10000.0) / (d_model // 2)))
+        if self.temp_bug_fix:
+            div_term = torch.exp(
+                torch.arange(0, self.d_model // 2, 2).float() * (-math.log(10000.0) / (self.d_model // 2))
+            )
         else:  # a buggy implementation (for backward compatibility only)
-            div_term = torch.exp(torch.arange(0, d_model // 2, 2).float() * (-math.log(10000.0) / d_model // 2))
+            div_term = torch.exp(
+                torch.arange(0, self.d_model // 2, 2).float() * (-math.log(10000.0) / self.d_model // 2)
+            )
         div_term = div_term[:, None, None]  # [C//4, 1, 1]
         pe[0::4, :, :] = torch.sin(x_position * div_term)
         pe[1::4, :, :] = torch.cos(x_position * div_term)
         pe[2::4, :, :] = torch.sin(y_position * div_term)
         pe[3::4, :, :] = torch.cos(y_position * div_term)
+        return pe.unsqueeze(0)
 
-        self.register_buffer('pe', pe.unsqueeze(0), persistent=False)  # [1, C, H, W]
+    def update_position_encoding_size(self, max_shape):
+        """Updates position encoding to new max_shape.
+
+        For 1/8 feature map (which is standard): If the input image size is H, W (both divisible by 8), the max_shape
+        should be (H//8, W//8).
+        """
+        self.pe = self._create_position_encoding(max_shape).to(self.pe.device)
 
     def forward(self, x):
         """
         Args:
             x: [N, C, H, W]
         """
+        if x.size(2) > self.pe.size(2) or x.size(3) > self.pe.size(3):
+            max_shape = (max(x.size(2), self.pe.size(2)), max(x.size(3), self.pe.size(3)))
+            self.update_position_encoding_size(max_shape)
+
         return x + self.pe[:, :, : x.size(2), : x.size(3)]


Test patch: