[Doc] Fix set_axis_separator example (#16792)

Minor fix to update the `set_axis_separator` example to match the definition
apache · Mar 26, 2024 · d43e1ab · d43e1ab
1 parent a768ee4
commit d43e1ab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py
@@ -3542,7 +3542,7 @@ def before_set_axis_separator(
         .. code-block:: python
 
             sch = tir.Schedule(before_set_axis_separator)
-            sch.set_axis_separators(sch.get_block("B"), buffer_index=0, buffer_index_type="write",
+            sch.set_axis_separators(sch.get_block("B"), buffer=("write", 0),
                                     axis_separators=[1])
             print(sch.mod["main"].script())