consider device number for xla

pytorch · Jan 22, 2024 · 047ea00 · 047ea00
1 parent a694479
commit 047ea00
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py
@@ -230,7 +230,7 @@ def load_benchmark(self):
 
     # torchbench uses `xla` as device instead of `tpu`
     if device := self.benchmark_experiment.accelerator == 'tpu':
-      device = 'xla'
+      device = str(self.benchmark_experiment.get_device())
     return benchmark_cls(
         test=self.benchmark_experiment.test,
         device=device,