diff --git a/yolov6/utils/figure_iou.py b/yolov6/utils/figure_iou.py
index 13b69d77..1d4b7589 100644
--- a/yolov6/utils/figure_iou.py
+++ b/yolov6/utils/figure_iou.py
@@ -83,6 +83,8 @@ def __call__(self, box1, box2):
             loss = loss.sum()
         elif self.reduction == 'mean':
             loss = loss.mean()
+            
+        loss = torch.nan_to_num(loss)
 
         return loss