!48302 fix segmentation fault raised by multimarginloss

Merge pull request !48302 from OwenSec/master
This commit is contained in:
i-robot 2023-02-02 01:13:32 +00:00 committed by Gitee
commit ab68927414
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
1 changed files with 2 additions and 2 deletions

View File

@ -134,13 +134,13 @@ string MultiMarginLoss::get_reduction() const {
AbstractBasePtr MultiMarginLossInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, AbstractBasePtr MultiMarginLossInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) { const std::vector<AbstractBasePtr> &input_args) {
MS_EXCEPTION_IF_NULL(primitive); MS_EXCEPTION_IF_NULL(primitive);
CheckAndConvertUtils::CheckInRange("multi_margin_loss_input_nums", input_args.size(), kIncludeBoth, {kDim2, kDim3},
primitive->name());
MS_EXCEPTION_IF_NULL(input_args[kInputIndex0]); MS_EXCEPTION_IF_NULL(input_args[kInputIndex0]);
MS_EXCEPTION_IF_NULL(input_args[kInputIndex1]); MS_EXCEPTION_IF_NULL(input_args[kInputIndex1]);
if (input_args.size() == kDim3) { if (input_args.size() == kDim3) {
MS_EXCEPTION_IF_NULL(input_args[kInputIndex2]); MS_EXCEPTION_IF_NULL(input_args[kInputIndex2]);
} }
CheckAndConvertUtils::CheckInRange("multi_margin_loss_input_nums", input_args.size(), kIncludeBoth, {kDim2, kDim3},
primitive->name());
auto types = MultiMarginLossInferType(primitive, input_args); auto types = MultiMarginLossInferType(primitive, input_args);
auto shapes = MultiMarginLossInferShape(primitive, input_args); auto shapes = MultiMarginLossInferShape(primitive, input_args);
return abstract::MakeAbstract(shapes, types); return abstract::MakeAbstract(shapes, types);