Попытка добавить onnx в работе с nvidia

src/networks/AMT-S.py

@@ -67,7 +67,14 @@ class Model(nn.Module):
         flow = torch.cat([flow0, flow1], dim=1)
         return corr, flow
 
-    def forward(self, img0, img1, embt, scale_factor=1.0, eval=False, **kwargs):
+    def forward(
+        self,
+        img0: torch.Tensor,
+        img1: torch.Tensor,
+        embt: torch.Tensor,
+    ):
+        scale_factor = 1.0
+        eval = False
         mean_ = (
             torch.cat([img0, img1], 2)
             .mean(1, keepdim=True)