fix end2end export

furyhawk · Sep 14, 2022 · 16301c0 · 16301c0
1 parent 21e0f28
commit 16301c0
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/yolov6/models/end2end.py b/yolov6/models/end2end.py
@@ -255,6 +255,9 @@ def forward(self, x):
         if self.with_preprocess:
             x = x[:,[2,1,0],...]
             x = x * (1/255)
-        x, _ = self.model(x)
+        if isinstance(x, list):
+            x = x[0]
+        else:
+            x = x
         x = self.end2end(x)
         return x