update

pytorch · Aug 8, 2024 · 67fbb06 · 67fbb06
1 parent e4e7a7b
commit 67fbb06
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torchao/quantization/autoquant.py b/torchao/quantization/autoquant.py
@@ -26,7 +26,7 @@
     try:
         from torch._inductor.runtime.runtime_utils import do_bench
     except ImportError:
-        from torch._inductor.runtime.benchmarking import do_bench
+        from torch._inductor.runtime.benchmarking.benchmarker import benchmark as do_bench
 
 
 __all__ = [
@@ -235,9 +235,8 @@ def do_autoquant_bench(op, *args, **kwargs):
         if TORCH_VERSION_AFTER_2_3 and not TORCH_VERSION_AFTER_2_5:
             from torch._inductor.runtime.runtime_utils import do_bench_gpu
             res = do_bench_gpu(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
-
         elif TORCH_VERSION_AFTER_2_5:
-            from torch._inductor.runtime.benchmarking import do_bench_gpu
+            from torch._inductor.runtime.benchmarking.benchmarker import benchmark_gpu as do_bench_gpu
             res = do_bench_gpu(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")
         else:
             res = do_bench(lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median")