From e9567ccb3f24adb766818ce7cbcc17a219ff0214 Mon Sep 17 00:00:00 2001
From: zhiqiang <zhiqwang@foxmail.com>
Date: Thu, 25 Nov 2021 17:05:53 +0800
Subject: [PATCH] Fixing jit trace

---
 yolort/models/box_head.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/yolort/models/box_head.py b/yolort/models/box_head.py
index 433df1fc..aa17f089 100644
--- 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)