[Fix] Fix nan loss in RotatedIoULoss when using AmpOptimizer (open-mm…

…lab#889) * [Fix] Fix nan loss in RotatedIoULoss when using AmpOptimizer * Fix lint
triple-Mu · Jul 4, 2023 · d50ab76 · d50ab76
1 parent 4aae1fc
commit d50ab76
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/mmrotate/models/losses/rotated_iou_loss.py b/mmrotate/models/losses/rotated_iou_loss.py
@@ -127,13 +127,14 @@ def forward(self,
             # iou_loss of shape (n,)
             assert weight.shape == pred.shape
             weight = weight.mean(-1)
-        loss = self.loss_weight * rotated_iou_loss(
-            pred,
-            target,
-            weight,
-            mode=self.mode,
-            eps=self.eps,
-            reduction=reduction,
-            avg_factor=avg_factor,
-            **kwargs)
+        with torch.cuda.amp.autocast(enabled=False):
+            loss = self.loss_weight * rotated_iou_loss(
+                pred,
+                target,
+                weight,
+                mode=self.mode,
+                eps=self.eps,
+                reduction=reduction,
+                avg_factor=avg_factor,
+                **kwargs)
         return loss