[TIR][FIX] check args size when creating prim_func by runtime::Regist…

…ry (apache#13809)
cblmemo · Jan 31, 2023 · 558c994 · 558c994
1 parent c81aaa8
commit 558c994
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
@@ -574,6 +574,7 @@ TVM_REGISTER_GLOBAL("te.CreatePrimFunc").set_body([](TVMArgs args, TVMRetValue*
   Array<te::Tensor> arg_list = args[0];
   std::optional<DataType> index_dtype_override{std::nullopt};
   // Add conversion to make std::optional compatible with FFI.
+  ICHECK_EQ(args.size(), 2);
   if (args[1].type_code() != kTVMNullptr) {
     index_dtype_override = args[1].operator DataType();
   }