Fix the potential index overflow (#3751)

diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc
index 07315d5..03a92b3 100644 (file)
--- a/src/relay/op/tensor/transform.cc
+++ b/src/relay/op/tensor/transform.cc
@@ -1451,9 +1451,10 @@ bool WhereRel(const Array<Type>& 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;