From b7d4305d24ec291c4aa4bdd84473cdfad8e21fc6 Mon Sep 17 00:00:00 2001 From: cchung100m Date: Sun, 11 Aug 2019 16:46:10 +0800 Subject: [PATCH] Fix the potential index overflow --- 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 07315d5d8098..03a92b35d396 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;