소스 검색

debug YOLOX-style Transform with Rotation

yjh0410 2 년 전
부모
커밋
8f6f8fa7dd
1개의 변경된 파일2개의 추가작업 그리고 2개의 파일을 삭제
  1. 2 2
      engine.py

+ 2 - 2
engine.py

@@ -209,7 +209,7 @@ class YoloTrainer(object):
             with torch.cuda.amp.autocast(enabled=self.args.fp16):
                 outputs = model(images)
                 # loss
-                loss_dict = self.criterion(outputs=outputs, targets=targets)
+                loss_dict = self.criterion(outputs=outputs, targets=targets, epoch=self.epoch)
                 losses = loss_dict['losses']
                 losses *= images.shape[0]  # loss * bs
 
@@ -786,7 +786,7 @@ class DetrTrainer(object):
             with torch.cuda.amp.autocast(enabled=self.args.fp16):
                 outputs = model(images)
                 # Compute loss
-                loss_dict = self.criterion(outputs=outputs, targets=targets)
+                loss_dict = self.criterion(outputs=outputs, targets=targets, epoch=self.epoch)
                 losses = loss_dict['losses']
 
                 loss_dict_reduced = distributed_utils.reduce_dict(loss_dict)