Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

add check of batchnorm args size

This commit is contained in:
tangxl 2023-02-09 20:39:29 +08:00
parent 9950d262c0
commit c0a3d589fd
2 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mindspore/core/ops/batch_norm.cc
View File

@ -119,6 +119,7 @@ class BatchNormInfer : public abstract::OpInferBase {
BaseShapePtr InferShape(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) const override {
const auto prim_name = primitive->name();
(void)CheckAndConvertUtils::CheckInputArgs(input_args, kGreaterThan, 0, prim_name);
auto x_shape_ptr = input_args[kInputIndex0]->BuildShape();
auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(x_shape_ptr)[kShape];
auto scale_shape_ptr = input_args[kInputIndex1]->BuildShape();
@ -175,6 +176,8 @@ class BatchNormInfer : public abstract::OpInferBase {
}
TypePtr InferType(const PrimitivePtr &prim, const std::vector<AbstractBasePtr> &input_args) const override {
const auto prim_name = prim->name();
(void)CheckAndConvertUtils::CheckInputArgs(input_args, kGreaterThan, 0, prim_name);
const std::set valid_types = {kFloat16, kFloat32};
auto x_type = input_args[0]->BuildType();
(void)CheckAndConvertUtils::CheckTensorTypeValid("input_x", x_type, valid_types, prim->name());

1
mindspore/core/ops/bn_infer.cc
View File

@ -14,7 +14,6 @@
* limitations under the License.
*/
#include <string>
#include <algorithm>
#include <memory>
#include <set>
#include <vector>