fix yolov4 EvalWhileTrain bug

mindspore-lab · Aug 22, 2023 · 7ac736b · 7ac736b
1 parent 07997fd
commit 7ac736b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mindyolo/models/heads/yolov4_head.py b/mindyolo/models/heads/yolov4_head.py
@@ -115,4 +115,7 @@ def construct(self, x):
             return prediction, box_xy, box_wh
         box_xy *= input_shape
         box_wh *= input_shape
-        return ops.concat((box_xy, box_wh, box_confidence, box_probs), -1)
+        return ops.concat((box_xy.astype(ms.float32),
+                           box_wh.astype(ms.float32),
+                           box_confidence.astype(ms.float32),
+                           box_probs.astype(ms.float32)), -1)