diff --git a/yolov6/models/end2end.py b/yolov6/models/end2end.py
index c1f102ba..e9f41b51 100644
--- a/yolov6/models/end2end.py
+++ b/yolov6/models/end2end.py
@@ -171,7 +171,7 @@ def forward(self, x):
         batched_dets = batched_dets.where((batch_inds == batch_template.unsqueeze(1)).unsqueeze(-1),batched_dets.new_zeros(1))
 
         batched_labels = cls_inds.unsqueeze(0).repeat(batch, 1)
-        batched_labels = batched_labels.where((batch_inds == batch_template.unsqueeze(1)),batched_labels.new_ones(1) * -1)
+        batched_labels = batched_labels.where((batch_inds == batch_template.unsqueeze(1)),batched_labels.new_ones(1, dtype=batched_labels.dtype) * -1)
 
         N = batched_dets.shape[0]
 
diff --git a/yolov6/models/losses/loss.py b/yolov6/models/losses/loss.py
index ec534923..86326965 100644
--- a/yolov6/models/losses/loss.py
+++ b/yolov6/models/losses/loss.py
@@ -205,7 +205,7 @@ def __init__(self):
     def forward(self, pred_score,gt_score, label, alpha=0.75, gamma=2.0):
 
         weight = alpha * pred_score.pow(gamma) * (1 - label) + gt_score * label
-        with torch.cuda.amp.autocast(enabled=False):
+        with torch.amp.autocast('cuda', enabled=False):
             loss = (F.binary_cross_entropy(pred_score.float(), gt_score.float(), reduction='none') * weight).sum()
 
         return loss