yjh0410 2 жил өмнө
parent
commit
cd1a7008cb

+ 1 - 1
models/yolov1/loss.py

@@ -34,7 +34,7 @@ class Criterion(object):
     
     
 
 
     def loss_labels(self, pred_cls, gt_label):
     def loss_labels(self, pred_cls, gt_label):
-        loss_cls = F.cross_entropy(pred_cls.permute(0, 2, 1), gt_label, reduction='none')
+        loss_cls = F.cross_entropy(pred_cls, gt_label, reduction='none')
 
 
         return loss_cls
         return loss_cls