diff --git a/tf_adapter/ops/aicpu/npu_cpu_ops.cc b/tf_adapter/ops/aicpu/npu_cpu_ops.cc
index 266b9d98843990a113d3d25c741d46f311c7762a..7bce51842666a11e0399d016167f18a828f66487 100644
--- a/tf_adapter/ops/aicpu/npu_cpu_ops.cc
+++ b/tf_adapter/ops/aicpu/npu_cpu_ops.cc
@@ -325,6 +325,7 @@ REGISTER_OP("FakeRemoteLookupUniqued")
   .Input("table_id: int32")
   .Input("keys: int64")
   .Input("actual_keys_input: int64")
+  .Input("key_count: int64")
   .Input("unique_indices: int32")
   .Output("values: float32")
   .Attr("embedding_dim: int = 0")