From 67a3c76a1298f6662edf948758729cb94eac25e5 Mon Sep 17 00:00:00 2001 From: YahooKID <40327829+YahooKID@users.noreply.github.com> Date: Tue, 12 Sep 2023 18:48:18 +0800 Subject: [PATCH 1/2] fix bug for inference with cpu --- yolov6/models/heads/effidehead_fuseab_seg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolov6/models/heads/effidehead_fuseab_seg.py b/yolov6/models/heads/effidehead_fuseab_seg.py index 80272928..bc4802b7 100644 --- a/yolov6/models/heads/effidehead_fuseab_seg.py +++ b/yolov6/models/heads/effidehead_fuseab_seg.py @@ -263,7 +263,7 @@ def forward(self, x): cls_output_af = torch.sigmoid(cls_output_af) # seg_output_af = torch.sigmoid(seg_output_af) - proto_no = (torch.ones(b, 1, l) * i).cuda() + proto_no = (torch.ones(b, 1, l) * i).to(device) if self.export: From 8b4f6b0b1d62a7ed5f809b80da97ceca74a3a779 Mon Sep 17 00:00:00 2001 From: YahooKID <40327829+YahooKID@users.noreply.github.com> Date: Tue, 12 Sep 2023 18:49:08 +0800 Subject: [PATCH 2/2] fix bug for inference with cpu --- yolov6/models/heads/effidehead_fuseab_seg_solo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolov6/models/heads/effidehead_fuseab_seg_solo.py b/yolov6/models/heads/effidehead_fuseab_seg_solo.py index 61bd1328..f1074be4 100644 --- a/yolov6/models/heads/effidehead_fuseab_seg_solo.py +++ b/yolov6/models/heads/effidehead_fuseab_seg_solo.py @@ -252,7 +252,7 @@ def forward(self, x): reg_output_af = self.proj_conv(F.softmax(reg_output_af, dim=1)) cls_output_af = torch.sigmoid(cls_output_af) - proto_no = (torch.ones(b, 1, l) * i).cuda() + proto_no = (torch.ones(b, 1, l) * i).to(device) if self.export: