Debug attempt

lanl · Jul 23, 2024 · 644c189 · 644c189
1 parent 5f7001e
commit 644c189
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/scico/flax/examples/ray_functions.py b/scico/flax/examples/ray_functions.py
@@ -163,8 +163,11 @@ def distributed_data_generation(
     @ray.remote(num_gpus=0.001)
     def data_gen(seed, size, ndata, imgf):
         import os
+        import sys
 
         os.environ["JAX_PLATFORMS"] = "cpu"
+        sys.modules.pop("jax")
+        sys.modules.pop("scico")
         os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
         os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
         if "CUDA_VISIBLE_DEVICES" in os.environ: