From 67fbb06a1988140e71f1f38ebc51a5d1ffbe64ea Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Thu, 8 Aug 2024 08:27:39 -0700 Subject: [PATCH] update --- torchao/quantization/autoquant.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/torchao/quantization/autoquant.py b/torchao/quantization/autoquant.py index 1e24701c3..c9caef4e2 100644 --- 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")