Fixing jit trace

zhiqwang · Nov 25, 2021 · e9567cc · e9567cc
1 parent a42e87b
commit e9567cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/yolort/models/box_head.py b/yolort/models/box_head.py
@@ -368,7 +368,7 @@ def forward(
                 shape of the element is (N, A, H, W, K).
             anchors_tuple (Tuple[Tensor, Tensor, Tensor]):
         """
-        batch_size = head_outputs[0][0]
+        batch_size = len(head_outputs[0])
 
         all_pred_logits = _concat_pred_logits(head_outputs)