fix bug in torch 1.13

681713de · Bubbliiiing · 1768ac04 · 681713de
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

nets/yolo_training.py nets/yolo_training.py +1 -1

未找到文件。
--- a/nets/yolo_training.py
+++ b/nets/yolo_training.py
@@ -384,7 +384,7 @@ class YOLOLoss(nn.Module):
            #-------------------------------------------#
            #   取出符合条件的框
            #-------------------------------------------#
-            from_which_layer    = from_which_layer[fg_mask_inboxes]
+            from_which_layer    = from_which_layer.to(fg_mask_inboxes.device)[fg_mask_inboxes]
            all_b               = all_b[fg_mask_inboxes]
            all_a               = all_a[fg_mask_inboxes]
            all_gj              = all_gj[fg_mask_inboxes]