From 6c5acd238c1597d4fbf589a6084ccc9f11d6b54a Mon Sep 17 00:00:00 2001 From: zhubochao <49826738+zhubochao@users.noreply.github.com> Date: Tue, 27 Dec 2022 11:54:28 +0800 Subject: [PATCH] Fix the device inconsistency error in yolov7 training (#397) * [FIX] Indices should be either on cpu or on the same device as the indexed tensor * fix lint error * use new_full instead of simple torch.ones * fix lint error * fix type error * use existed tensor to create new_full tensor Co-authored-by: zhubochao --- mmyolo/models/task_modules/assigners/batch_yolov7_assigner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mmyolo/models/task_modules/assigners/batch_yolov7_assigner.py b/mmyolo/models/task_modules/assigners/batch_yolov7_assigner.py index d955df97..0b330f99 100644 --- a/mmyolo/models/task_modules/assigners/batch_yolov7_assigner.py +++ b/mmyolo/models/task_modules/assigners/batch_yolov7_assigner.py @@ -227,7 +227,8 @@ class BatchYOLOv7Assigner(nn.Module): _mlvl_priors.append(priors) _from_which_layer.append( - torch.ones(size=(_mlvl_positive_info.shape[0], )) * i) + _mlvl_positive_info.new_full( + size=(_mlvl_positive_info.shape[0], ), fill_value=i)) # (n,85) level_batch_idx, prior_ind, \