[benchmarks] Fix CUDA device check. (pytorch#6419)

amithrm · Mar 1, 2024 · 23e36f8 · 23e36f8
1 parent 8221cb4
commit 23e36f8
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/benchmarks/benchmark_experiment.py b/benchmarks/benchmark_experiment.py
@@ -96,10 +96,7 @@ def _is_available(self, experiment_config):
     if cfg_accelerator == "tpu":
       if cfg_xla is None:
         return False
-    elif cfg_accelerator == "cuda":
-      if cfg_xla is None and not is_xla_device_available("CUDA"):
-        return False
-    elif cfg_accelerator == "cpu":
+    elif cfg_accelerator in ("cpu", "cuda"):
       pass
     else:
       raise NotImplementedError