Merge pull request #14502 from branfosj/20211204121131_new_pr_jax0219

pass tests on A* GPUs by setting 'NVIDIA_TF32_OVERRIDE=0' in jax 0.2.19
easybuilders · Dec 8, 2021 · 922fb7b · 922fb7b
2 parents 3dc63e2 + a1fc786
commit 922fb7b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/easybuild/easyconfigs/j/jax/jax-0.2.19-fosscuda-2020b.eb b/easybuild/easyconfigs/j/jax/jax-0.2.19-fosscuda-2020b.eb
@@ -115,7 +115,8 @@ exts_list = [
         # see https://github.com/google/jax/issues/7323 and
         # https://github.com/google/jax/blob/main/docs/gpu_memory_allocation.rst;
         # use CUDA_VISIBLE_DEVICES=0 to avoid failing tests on systems with multiple GPUs;
-        'runtest': 'CUDA_VISIBLE_DEVICES=0 JAX_ENABLE_X64=true pytest tests',
+        # use NVIDIA_TF32_OVERRIDE=0 to avoid lossing numerical precision by disabling TF32 Tensor Cores;
+        'runtest': "NVIDIA_TF32_OVERRIDE=0 CUDA_VISIBLE_DEVICES=0 JAX_ENABLE_X64=true pytest tests",
     }),
 ]