fix bug for inference with cpu

meituan · Sep 12, 2023 · 8b4f6b0 · 8b4f6b0
1 parent 67a3c76
commit 8b4f6b0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: