gpu floatstatus add type check

5 years ago · 2b51199054
parent f23bfe0d71
commit 2b51199054
1 changed files with 1 additions and 0 deletions
--- a/mindspore/ops/operations/math_ops.py
+++ b/mindspore/ops/operations/math_ops.py
@ -1712,6 +1712,7 @@ class FloatStatus(PrimitiveWithInfer):
        return [1]

    def infer_dtype(self, x_dtype):
+        validator.check_tensor_type_same({'x': x_dtype}, [mstype.float32, mstype.float16], self.name)
        return x_dtype

 class NPUAllocFloatStatus(PrimitiveWithInfer):