fix end2end export

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