Add a check for multiplex op (PaddlePaddle#34972)

esythan · Sep 26, 2021 · b430f6a · b430f6a
1 parent 628ff34
commit b430f6a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/paddle/fluid/operators/multiplex_op.cu b/paddle/fluid/operators/multiplex_op.cu
@@ -29,6 +29,14 @@ class MultiplexGPUKernel : public framework::OpKernel<T> {
     auto* out = ctx.Output<Tensor>("Out");
     out->mutable_data<T>(ctx.GetPlace());
 
+    for (size_t i = 0; i < ins.size(); ++i) {
+      PADDLE_ENFORCE_GT(
+          ins[i]->numel(), 0,
+          platform::errors::OutOfRange(
+              "indexing will be out of bounds with size 0 for the %d-th input.",
+              i));
+    }
+
     auto rows = ins[0]->dims()[0];
     auto cols = ins[0]->numel() / rows;
     // copy index to cpu

diff --git a/paddle/fluid/operators/multiplex_op.h b/paddle/fluid/operators/multiplex_op.h
@@ -31,6 +31,14 @@ class MultiplexCPUKernel : public framework::OpKernel<T> {
 
     out->mutable_data<T>(ctx.GetPlace());
 
+    for (size_t i = 0; i < ins.size(); ++i) {
+      PADDLE_ENFORCE_GT(
+          ins[i]->numel(), 0,
+          platform::errors::OutOfRange(
+              "indexing will be out of bounds with size 0 for the %d-th input.",
+              i));
+    }
+
     auto rows = ins[0]->dims()[0];
     auto cols = ins[0]->numel() / rows;
     auto index = ids->data<int32_t>();