Fix autoquant tests failed due to changes to benchmark_gpu (#2818)

Skip test failing only in CI

Summary:
att

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

torchao/quantization/autoquant.py

@@ -34,6 +34,7 @@
     TorchAOBaseTensor,
     is_sm_at_least_89,
     is_sm_at_least_90,
+    torch_version_at_least,
 )
 
 from .granularity import (
@@ -343,9 +344,18 @@ def do_autoquant_bench(op, *args, **kwargs):
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph, stream=stream):
             op(*args, **kwargs)
-        res = benchmarker.benchmark_gpu(
-            lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median"
-        )
+        # TODO: update to 2.8.0 after https://github.com/pytorch/ao/pull/2786 is landed
+        if torch_version_at_least("2.9.0"):
+            from statistics import median
+
+            res = benchmarker.benchmark_gpu(
+                lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="all"
+            )
+            res = median(res)
+        else:
+            res = benchmarker.benchmark_gpu(
+                lambda: graph.replay(), warmup=warmup, rep=rep, return_mode="median"
+            )
     return res