diff --git a/Golden_Repo/j/JAX/JAX-0.2.9-gcccoremkl-9.3.0-2020.2.254-Python-3.8.5.eb b/Golden_Repo/j/JAX/JAX-0.2.9-gcccoremkl-9.3.0-2020.2.254-Python-3.8.5.eb
index 217069da31ac73c3df869cad6e7ca7d76b65ae88..eba72dd937ddf100fd689bd9c65b19c5c7c9968d 100644
--- a/Golden_Repo/j/JAX/JAX-0.2.9-gcccoremkl-9.3.0-2020.2.254-Python-3.8.5.eb
+++ b/Golden_Repo/j/JAX/JAX-0.2.9-gcccoremkl-9.3.0-2020.2.254-Python-3.8.5.eb
@@ -1,3 +1,4 @@
+import os
 easyblock = 'PythonPackage'
 
 name = 'JAX'
@@ -14,6 +15,7 @@ site_contacts = 'a.strube@fz-juelich.de'
 
 dependencies = [
     ('binutils', '2.34'),
+    ('CUDA', '11.0', '', SYSTEM),
     ('Python', '3.8.5'),
     ('SciPy-Stack', '2020', versionsuffix, ('gcccoremkl', '9.3.0-2020.2.254')),
     ('cuDNN', '8.0.2.39', '-CUDA-%s' % local_cudaver, True),
@@ -56,10 +58,11 @@ exts_list = [
     }),
 ]
 
-
-modextravars = {
-    'XLA_FLAGS': '--xla_gpu_cuda_data_dir=$EBROOTCUDA',
-}
+# This should be modextravars, but life is unfair
+modluafooter = """
+setenv("CUDA_DIR", os.getenv("EBROOTCUDA"))
+setenv("XLA_FLAGS", "--xla_gpu_cuda_data_dir=" .. os.getenv("EBROOTCUDA"))
+"""
 
 sanity_check_paths = {
     'files': [],