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:
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: