Update benchmarks/benchmark.py

benchmarks/benchmark.py

@@ -4,11 +4,20 @@ import torch.nn.functional as F
 from kernels.benchmark import Benchmark
 
 
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
 class ReluBenchmark(Benchmark):
     seed: int = 42
 
     def setup(self):
-        self.x = torch.randn(1024, 1024, device="cuda", dtype=torch.float32)
+        device = get_device()
+        self.x = torch.randn(1024, 1024, device=device, dtype=torch.float32)
         self.out = torch.empty_like(self.x)
 
     def benchmark_base(self):
@@ -18,7 +27,8 @@ class ReluBenchmark(Benchmark):
         return F.relu(self.x)
 
     def setup_large(self):
-        self.x = torch.randn(4096, 4096, device="cuda", dtype=torch.float32)
+        device = get_device()
+        self.x = torch.randn(4096, 4096, device=device, dtype=torch.float32)
         self.out = torch.empty_like(self.x)
 
     def benchmark_large(self):