diff --git a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py
index ab93d150c..7f64c89a8 100644
--- a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py
+++ b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py
@@ -20,14 +20,17 @@
 using standard environment variables like JAX_COORDINATOR_ADDRESS,
 JAX_PROCESS_ID and JAX_NUM_PROCESSES.
 """
-# pylint: disable=g-statement-before-imports,g-import-not-at-top
+import os
 
-try:  # SimDevice import must occur before JAX.
-  import simdevice  # pylint: disable=unused-import
-except (ImportError, FileNotFoundError):
-  pass
+# pylint: disable=g-statement-before-imports,g-import-not-at-top
+if os.environ.get('SIMDEVICE_CONFIG_PATH') or os.environ.get(
+    'ENABLE_SIMDEVICE', ''
+).lower() in ('true', '1'):
+  try:  # SimDevice import must occur before JAX.
+    import simdevice  # pylint: disable=unused-import
+  except (ImportError, FileNotFoundError):
+    pass
 
-import os
 
 from absl import app
 from absl import flags