Update benchmarks/benchmark.py

benchmarks/benchmark.py

@@ -4,20 +4,11 @@ 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):
-        device = get_device()
-        self.x = torch.randn(1024, 1024, device=device, dtype=torch.float32)
+        self.x = torch.randn(1024, 1024, device=self.device, dtype=torch.float32)
         self.out = torch.empty_like(self.x)
 
     def benchmark_base(self):
@@ -27,8 +18,7 @@ class ReluBenchmark(Benchmark):
         return F.relu(self.x)
 
     def setup_large(self):
-        device = get_device()
-        self.x = torch.randn(4096, 4096, device=device, dtype=torch.float32)
+        self.x = torch.randn(4096, 4096, device=self.device, dtype=torch.float32)
         self.out = torch.empty_like(self.x)
 
     def benchmark_large(self):