From 23e36f84c5594e2fdda8203e561e050059580211 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Wed, 31 Jan 2024 10:42:25 -0300 Subject: [PATCH] [benchmarks] Fix CUDA device check. (#6419) --- benchmarks/benchmark_experiment.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/benchmarks/benchmark_experiment.py b/benchmarks/benchmark_experiment.py index 2b48c7e03de4..187acaff3ded 100644 --- 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