diff --git a/mindspore/core/ops/op_utils.h b/mindspore/core/ops/op_utils.h index 5d641628064..1a180ed807a 100644 --- a/mindspore/core/ops/op_utils.h +++ b/mindspore/core/ops/op_utils.h @@ -29,6 +29,9 @@ namespace mindspore::ops { const std::set common_valid_types = {kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16, kUInt32, kUInt64, kFloat16, kFloat32, kFloat64}; +const std::set common_valid_types_with_bool = {kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16, + kUInt32, kUInt64, kFloat16, kFloat32, kFloat64, kBool}; + const std::set common_valid_types_with_complex = {kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16, kUInt32, kUInt64, kFloat16, kFloat32, kFloat64, kComplex64, kComplex128}; diff --git a/mindspore/core/ops/tensor_scatter_arithmetic.cc b/mindspore/core/ops/tensor_scatter_arithmetic.cc index 919a1392a0a..fe183ce9c52 100644 --- a/mindspore/core/ops/tensor_scatter_arithmetic.cc +++ b/mindspore/core/ops/tensor_scatter_arithmetic.cc @@ -83,6 +83,9 @@ TypePtr TensorScatterArithmeticInferType(const PrimitivePtr &primitive, std::map type_dict; type_dict.emplace("input_x", input_args[kInputIndex0]->BuildType()); type_dict.emplace("updates", input_args[kInputIndex2]->BuildType()); + if (prim_name == prim::kPrimTensorScatterUpdate->name()) { + return CheckAndConvertUtils::CheckTensorTypeSame(type_dict, common_valid_types_with_bool, prim_name); + } return CheckAndConvertUtils::CheckTensorTypeSame(type_dict, common_valid_types, prim_name); } } // namespace