Add pruno0,1 and stop into docstring

cvg · Jan 26, 2024 · 5046dbe · 5046dbe
1 parent 3d5e4d0
commit 5046dbe
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/lightglue/lightglue.py b/lightglue/lightglue.py
@@ -463,7 +463,11 @@ def forward(self, data: dict) -> dict:
             matching_scores0: [B x M]
             matches1: [B x N]
             matching_scores1: [B x N]
-            matches: List[[Si x 2]], scores: List[[Si]]
+            matches: List[[Si x 2]],
+            scores: List[[Si]]
+            stop: int,
+            prune0: [B x M],
+            prune1: [B x N] 
         """
         with torch.autocast(enabled=self.conf.mp, device_type="cuda"):
             return self._forward(data)