From 9a1f858ee9f6fe19eee24bef1659a1c90aba98e7 Mon Sep 17 00:00:00 2001
From: Bubbliiiing <3323290568@qq.com>
Date: Sat, 1 Oct 2022 16:13:13 +0800
Subject: [PATCH] fix __int64 bug

---
 nets/yolo_training.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/nets/yolo_training.py b/nets/yolo_training.py
index d965e42..f59278c 100644
--- a/nets/yolo_training.py
+++ b/nets/yolo_training.py
@@ -448,6 +448,7 @@ class YOLOLoss(nn.Module):
             #   anchors_i [num_anchor, 2]
             #----------------------------------------------------#
             anchors_i = torch.from_numpy(self.anchors[i] / self.stride[i]).type_as(predictions[i])
+            anchors_i, shape = torch.from_numpy(self.anchors[i] / self.stride[i]).type_as(predictions[i]), predictions[i].shape
             #-------------------------------------------#
             #   计算获得对应特征层的高宽
             #-------------------------------------------#
@@ -509,7 +510,7 @@ class YOLOLoss(nn.Module):
             #   a代表属于该特征点的第几个先验框
             #-------------------------------------------#
             a = t[:, 6].long()  # anchor indices
-            indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
+            indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid indices
             anchors.append(anchors_i[a])  # anchors
 
         return indices, anchors
-- 
GitLab