From 5dd3d96d830fbbd6e1f28bb42c1182d7795ca7b2 Mon Sep 17 00:00:00 2001
From: Yiheng <mayiheng8@gmail.com>
Date: Thu, 19 Jan 2023 16:26:12 +0800
Subject: [PATCH] [TIR][FIX] check args size when creating prim_func by
 runtime::Registry

---
 src/te/operation/create_primfunc.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index 92186a4ffea4..bdcdbc023a1c 100644
--- 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();
   }