Improve consistency with similar examples

lanl · Oct 8, 2024 · 28b0a61 · 28b0a61
1 parent 68ff3b5
commit 28b0a61
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/examples/scripts/ct_datagen_foam2.py b/examples/scripts/ct_datagen_foam2.py
@@ -14,6 +14,7 @@
 """
 
 # isort: off
+import os
 import numpy as np
 
 import logging
@@ -24,6 +25,9 @@
         logging_level=logging.ERROR
     )  # need to call init before jax import: ray-project/ray#44087
 
+# Set an arbitrary processor count (only applies if GPU is not available).
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
+
 from scico import plot
 from scico.flax.examples import load_ct_data