From 9161efbc1f5769ba97beae11b74395206b741619 Mon Sep 17 00:00:00 2001 From: Neo Chien Date: Mon, 12 Aug 2019 23:54:18 +0800 Subject: [PATCH] Fix the potential index overflow (#3751) --- src/relay/op/tensor/transform.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 07315d5..03a92b3 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1451,9 +1451,10 @@ bool WhereRel(const Array& types, CHECK(reporter->AssertEQ(x_shape[i], y_shape[i])) << "x and y must have the same shape: " << x_shape << " vs " << y_shape; - CHECK(reporter->AssertEQ(cond_shape[i], x_shape[i])) - << "Shape of condition " << condition->shape - << " must be either equal to x or has dimension of 1."; + if (i < cond_shape.size()) { + CHECK(reporter->AssertEQ(cond_shape[i], x_shape[i])) + << "condition and x must have the same shape: " << cond_shape << " vs " << x_shape; + } } reporter->Assign(types[3], TensorTypeNode::make(x_shape, x->dtype)); return true; -- 2.7.4