From 047ea0008a360f9b9465c21a9c0e316c21118fab Mon Sep 17 00:00:00 2001
From: zpcore <zpcore@gmail.com>
Date: Mon, 22 Jan 2024 05:36:12 +0000
Subject: [PATCH] consider device number for xla

---
 benchmarks/torchbench_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py
index a3aa5185c6e..bd95f780b00 100644
--- 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,