Fix a bug in test core.

zheng-da · Dec 9, 2017 · 53eec60 · 53eec60
1 parent 7748180
commit 53eec60
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/cpp/include/test_core_op.h b/tests/cpp/include/test_core_op.h
@@ -314,7 +314,9 @@ class CoreOpExecutor : public test::op::OperatorDataInitializer<DType>
       // Set up forward
       attrs_ = ParseAttrs(op_, args);
 
-      const int num_inputs = op_->num_inputs;
+      int num_inputs = op_->num_inputs;
+      if (op_->get_num_inputs)
+        num_inputs = op_->get_num_inputs(attrs_);
 
       if (!inputs.empty()) {
         CHECK_EQ(inputs.size(), static_cast<size_t>(num_inputs));