Update __init__.py

unslothai · Jan 23, 2025 · 436fe17 · 436fe17
1 parent f291a4c
commit 436fe17
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/unsloth/__init__.py b/unsloth/__init__.py
@@ -153,8 +153,8 @@ def fix_triton_ops():
             early_config_prune,
             estimate_matmul_time,
         )
-        triton.ops.early_config_prune   = early_config_prune
-        triton.ops.estimate_matmul_time = estimate_matmul_time
+        triton.ops.matmul_perf_model.early_config_prune   = early_config_prune
+        triton.ops.matmul_perf_model.estimate_matmul_time = estimate_matmul_time
     pass
 pass
 fix_triton_ops()