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

meituan · Sep 15, 2023 · 463d8e7 · 463d8e7
2 parents c92a7ed + 8b4f6b0
commit 463d8e7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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:

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: