1 files changed, 8 insertions, 8 deletions
diff --git a/candle-core/benches/matmul.rs b/candle-core/benches/matmul.rs
index 83679771..4f7dfa6c 100644
--- a/candle-core/benches/matmul.rs
+++ b/candle-core/benches/matmul.rs
@@ -1,4 +1,8 @@
-use candle_core::{DType, Device, Tensor};
+mod bench_utils;
+
+use crate::bench_utils::bench_name;
+use bench_utils::{device, BenchDevice};
+use candle_core::{DType, Tensor};
 use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
 use std::time::Instant;
 
@@ -12,14 +16,14 @@ fn criterion_benchmark(c: &mut Criterion) {
     let n = 2048;
     let k = 2048;
 
-    let device = Device::new_metal(0).unwrap();
+    let device = device().unwrap();
     let dtype = DType::F32;
     let lhs = Tensor::zeros((b, m, k), dtype, &device).unwrap();
     let rhs = Tensor::zeros((b, n, k), dtype, &device).unwrap();
 
     let flops = b * m * n * k;
 
-    let mut group = c.benchmark_group("matmul_metal");
+    let mut group = c.benchmark_group(bench_name("matmul"));
     group.throughput(Throughput::Bytes(flops as u64));
     group.bench_function("iter", move |b| {
         b.iter_custom(|iters| {
@@ -27,11 +31,7 @@ fn criterion_benchmark(c: &mut Criterion) {
             for _i in 0..iters {
                 run(black_box(&lhs), black_box(&rhs));
             }
-            if let Device::Metal(device) = &device {
-                device.wait_until_completed().unwrap();
-            } else {
-                panic!("Expected metal device");
-            }
+            device.sync().unwrap();
             start.elapsed()
         })
     });