Merge branch 'yolov6-seg' of github.com:meituan/YOLOv6 into yolov6-seg

Author: Chilicyy

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:

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: