diff --git a/byte_micro_perf/core/ops/reduction_ops.py b/byte_micro_perf/core/ops/reduction_ops.py
index fac19b5f..bcb629d7 100644
--- a/byte_micro_perf/core/ops/reduction_ops.py
+++ b/byte_micro_perf/core/ops/reduction_ops.py
@@ -309,7 +309,7 @@ def prepare(self):
         self.io_bytes = self.read_bytes + self.write_bytes
 
         self._run_func = self.softmax_run
-        self._create_in_out_tensors_func = partial(
+        self._create_tensors_func = partial(
             self._create_in_out_tensors,
             create_inputs=True,
             create_outputs=False
@@ -365,7 +365,7 @@ def prepare(self):
         self.io_bytes = self.read_bytes + self.write_bytes
 
         self._run_func = self.topk_run
-        self._create_in_out_tensors_func = partial(
+        self._create_tensors_func = partial(
             self._create_in_out_tensors,
             create_inputs=True,
             create_outputs=False