RuntimeError: jaxlib/gpu/solver_kernels.cc:45: operation gpusolverDnCreate(&handle) failed: cuSolver internal error

**Description**
I tried to reproduce the results from the paper by running `benchmark1.py`, and it worked on CPU while encountering the error `cuSolver internal error` on GPU.

**Error**

```
Traceback (most recent call last):
File "/home//robust-reinforcement-learning/dmc_vision_benchmark-main/example/benchmark1.py", line 135, in
app.run(run_benchmark1)
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/absl/app.py", line 367, in run
_run_main(main, args)
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/absl/app.py", line 312, in _run_main
sys.exit(main(argv))
File "/home//robust-reinforcement-learning/dmc_vision_benchmark-main/example/benchmark1.py", line 108, in run_benchmark1
_, metrics = trainer.train()
File "/home//robust-reinforcement-learning/dmc_vision_benchmark-main/dmc_vision_benchmark/rep_learn/train.py", line 688, in train
model_params = self.model.init(
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/kauldron/typing/type_check.py", line 120, in _reraise_with_shape_info
retval = fn(*args, **kwargs)
File "/home//robust-reinforcement-learning/dmc_vision_benchmark-main/dmc_vision_benchmark/rep_learn/mlp_multiheads.py", line 72, in __call__
obs_encoded = self.obs_encoder(obs)
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/flax/linen/combinators.py", line 112, in __call__
outputs = layer(outputs)
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/kauldron/typing/type_check.py", line 120, in _reraise_with_shape_info
retval = fn(*args, **kwargs)
File "/home//robust-reinforcement-learning/dmc_vision_benchmark-main/dmc_vision_benchmark/modules/cnn.py", line 81, in __call__
x = nn.Conv(
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/flax/linen/linear.py", line 636, in __call__
kernel = self.param(
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/jax/_src/nn/initializers.py", line 611, in init
Q, R = jnp.linalg.qr(A)
File "/home//.conda/envs/jax-gpu/lib/python3.11/site-packages/jax/_src/numpy/linalg.py", line 1300, in qr
q, r = lax_linalg.qr(a, full_matrices=full_matrices)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: RuntimeError: jaxlib/gpu/solver_handle_pool.cc:37: operation gpusolverDnCreate(&handle) failed: cuSolver internal error

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
```

**Environment**

`Python 3.11.14`

`NVIDIA-SMI 580.126.09 Driver Version: 580.126.09 CUDA Version: 13.0 `

```
# Name Version Build Channel
_libgcc_mutex 0.1 main
_openmp_mutex 5.1 1_gnu
absl-py 2.4.0 pypi_0 pypi
altair 6.0.0 pypi_0 pypi
array-record 0.8.3 pypi_0 pypi
asttokens 3.0.1 pypi_0 pypi
astunparse 1.6.3 pypi_0 pypi
attrs 25.4.0 pypi_0 pypi
bzip2 1.0.8 h5eee18b_6
ca-certificates 2025.12.2 h06a4308_0
certifi 2026.1.4 pypi_0 pypi
charset-normalizer 3.4.4 pypi_0 pypi
chex 0.1.90 pypi_0 pypi
clu 0.0.12 pypi_0 pypi
contextlib2 21.6.0 pypi_0 pypi
contourpy 1.3.3 pypi_0 pypi
cycler 0.12.1 pypi_0 pypi
decorator 5.2.1 pypi_0 pypi
distracting-control 2025.6.3.7.25.22 pypi_0 pypi
dm-control 1.0.21.post1 pypi_0 pypi
dm-env 1.6 pypi_0 pypi
dm-reverb-nightly 0.15.0.dev20240214 pypi_0 pypi
dm-tree 0.1.9 pypi_0 pypi
dmc-vision-benchmark 0.0.1 pypi_0 pypi
docstring-parser 0.17.0 pypi_0 pypi
einops 0.8.2 pypi_0 pypi
etils 1.13.0 pypi_0 pypi
executing 2.2.1 pypi_0 pypi
expat 2.7.3 h7354ed3_4
flatbuffers 25.12.19 pypi_0 pypi
flax 0.8.5 pypi_0 pypi
fonttools 4.61.1 pypi_0 pypi
fsspec 2026.1.0 pypi_0 pypi
gast 0.7.0 pypi_0 pypi
glfw 2.10.0 pypi_0 pypi
google-pasta 0.2.0 pypi_0 pypi
googleapis-common-protos 1.72.0 pypi_0 pypi
grpcio 1.76.0 pypi_0 pypi
h5py 3.15.1 pypi_0 pypi
idna 3.11 pypi_0 pypi
immutabledict 4.2.2 pypi_0 pypi
importlib-resources 6.5.2 pypi_0 pypi
ipython 9.9.0 pypi_0 pypi
ipython-pygments-lexers 1.1.1 pypi_0 pypi
jax 0.4.34 pypi_0 pypi
jax-cuda12-pjrt 0.4.34 pypi_0 pypi
jax-cuda12-plugin 0.4.34 pypi_0 pypi
jaxlib 0.4.34 pypi_0 pypi
jaxtyping 0.3.6 pypi_0 pypi
jedi 0.19.2 pypi_0 pypi
jinja2 3.1.6 pypi_0 pypi
joblib 1.5.3 pypi_0 pypi
jsonschema 4.26.0 pypi_0 pypi
jsonschema-specifications 2025.9.1 pypi_0 pypi
kauldron 0.1.0 pypi_0 pypi
keras 3.13.1 pypi_0 pypi
kiwisolver 1.4.9 pypi_0 pypi
labmaze 1.0.6 pypi_0 pypi
lark 1.3.1 pypi_0 pypi
ld_impl_linux-64 2.44 h153f514_2
libclang 18.1.1 pypi_0 pypi
libexpat 2.7.3 h7354ed3_4
libffi 3.4.4 h6a678d5_1
libgcc 15.2.0 h69a1729_7
libgcc-ng 15.2.0 h166f726_7
libgomp 15.2.0 h4751f2c_7
libnsl 2.0.0 h5eee18b_0
libstdcxx 15.2.0 h39759b7_7
libstdcxx-ng 15.2.0 hc03a8fd_7
libuuid 1.41.5 h5eee18b_0
libxcb 1.17.0 h9b100fa_0
libzlib 1.3.1 hb25bd0a_0
lxml 6.0.2 pypi_0 pypi
markdown 3.10.1 pypi_0 pypi
markdown-it-py 4.0.0 pypi_0 pypi
markupsafe 3.0.3 pypi_0 pypi
matplotlib 3.10.8 pypi_0 pypi
matplotlib-inline 0.2.1 pypi_0 pypi
mdurl 0.1.2 pypi_0 pypi
mediapy 1.2.5 pypi_0 pypi
ml-collections 0.1.1 pypi_0 pypi
ml-dtypes 0.3.2 pypi_0 pypi
msgpack 1.1.2 pypi_0 pypi
mujoco 3.2.0 pypi_0 pypi
namex 0.1.0 pypi_0 pypi
narwhals 2.15.0 pypi_0 pypi
ncurses 6.5 h7934f7d_0
nest-asyncio 1.6.0 pypi_0 pypi
numpy 1.26.4 pypi_0 pypi
nvidia-cublas-cu12 12.9.1.4 pypi_0 pypi
nvidia-cuda-cupti-cu12 12.9.79 pypi_0 pypi
nvidia-cuda-nvcc-cu12 12.9.86 pypi_0 pypi
nvidia-cuda-nvrtc-cu12 12.9.86 pypi_0 pypi
nvidia-cuda-runtime-cu12 12.9.79 pypi_0 pypi
nvidia-cudnn-cu12 9.18.1.3 pypi_0 pypi
nvidia-cufft-cu12 11.4.1.4 pypi_0 pypi
nvidia-cusolver-cu12 11.7.5.82 pypi_0 pypi
nvidia-cusparse-cu12 12.5.10.65 pypi_0 pypi
nvidia-nccl-cu12 2.29.2 pypi_0 pypi
nvidia-nvjitlink-cu12 12.9.86 pypi_0 pypi
openssl 3.0.18 hd6dcaed_0
opt-einsum 3.4.0 pypi_0 pypi
optax 0.2.5 pypi_0 pypi
optree 0.18.0 pypi_0 pypi
orbax-checkpoint 0.5.23 pypi_0 pypi
packaging 25.0 py311h06a4308_1
parso 0.8.5 pypi_0 pypi
pexpect 4.9.0 pypi_0 pypi
pillow 12.1.0 pypi_0 pypi
pip 25.3 pyhc872135_0
portpicker 1.6.0 pypi_0 pypi
promise 2.3 pypi_0 pypi
prompt-toolkit 3.0.52 pypi_0 pypi
protobuf 4.25.8 pypi_0 pypi
psutil 7.2.1 pypi_0 pypi
pthread-stubs 0.3 h0ce48e5_1
ptyprocess 0.7.0 pypi_0 pypi
pure-eval 0.2.3 pypi_0 pypi
pyarrow 23.0.0 pypi_0 pypi
pygments 2.19.2 pypi_0 pypi
pyopengl 3.1.10 pypi_0 pypi
pyparsing 3.3.2 pypi_0 pypi
python 3.11.14 h6fa692b_0
python-dateutil 2.9.0.post0 pypi_0 pypi
pyyaml 6.0.3 pypi_0 pypi
readline 8.3 hc2a1206_0
referencing 0.37.0 pypi_0 pypi
regex 2026.1.15 pypi_0 pypi
requests 2.32.5 pypi_0 pypi
rich 14.3.1 pypi_0 pypi
rlds 0.1.8 pypi_0 pypi
rpds-py 0.30.0 pypi_0 pypi
scikit-fmm 2024.5.29 pypi_0 pypi
scikit-learn 1.8.0 pypi_0 pypi
scipy 1.17.0 pypi_0 pypi
setuptools 80.10.1 py311h06a4308_0
simple-parsing 0.1.8 pypi_0 pypi
six 1.17.0 pypi_0 pypi
sqlite 3.51.1 he0a8d7e_0
stack-data 0.6.3 pypi_0 pypi
tensorboard 2.16.2 pypi_0 pypi
tensorboard-data-server 0.7.2 pypi_0 pypi
tensorflow 2.16.2 pypi_0 pypi
tensorflow-datasets 4.9.9 pypi_0 pypi
tensorflow-io-gcs-filesystem 0.37.1 pypi_0 pypi
tensorflow-metadata 1.17.3 pypi_0 pypi
tensorstore 0.1.74 pypi_0 pypi
termcolor 3.3.0 pypi_0 pypi
threadpoolctl 3.6.0 pypi_0 pypi
tk 8.6.15 h54e0aa7_0
toml 0.10.2 pypi_0 pypi
toolz 1.1.0 pypi_0 pypi
tqdm 4.67.1 pypi_0 pypi
traitlets 5.14.3 pypi_0 pypi
typeguard 2.13.3.post54 pypi_0 pypi
typing-extensions 4.15.0 pypi_0 pypi
tzdata 2025c he532380_0
urllib3 2.6.3 pypi_0 pypi
wadler-lindig 0.1.7 pypi_0 pypi
wcwidth 0.5.0 pypi_0 pypi
werkzeug 3.1.5 pypi_0 pypi
wheel 0.46.3 py311h06a4308_0
wrapt 2.0.1 pypi_0 pypi
xorg-libx11 1.8.12 h9b100fa_1
xorg-libxau 1.0.12 h9b100fa_0
xorg-libxdmcp 1.1.5 h9b100fa_0
xorg-xorgproto 2024.1 h5eee18b_1
xz 5.6.4 h5eee18b_1
zipp 3.23.0 pypi_0 pypi
zlib 1.3.1 hb25bd0a_0
```

```
nvidia-cublas-cu12==12.9.1.4
nvidia-cuda-cupti-cu12==12.9.79
nvidia-cuda-nvcc-cu12==12.9.86
nvidia-cuda-nvrtc-cu12==12.9.86
nvidia-cuda-runtime-cu12==12.9.79
nvidia-cudnn-cu12==9.18.1.3
nvidia-cufft-cu12==11.4.1.4
nvidia-cusolver-cu12==11.7.5.82
nvidia-cusparse-cu12==12.5.10.65
nvidia-nccl-cu12==2.29.2
nvidia-nvjitlink-cu12==12.9.86
```

**Library Installation**

`pip install "jax[cuda12]==0.4.34" "jaxlib==0.4.34" "jax-cuda12-plugin==0.4.34" "jax-cuda12-pjrt==0.4.34"`

**Run**

```
#!/bin/bash
#SBATCH --job-name=dmc-benchmark-1 # Name of your job
#SBATCH --output=logs/%x_%j.out # Output file: logs/jobname_jobid.out
#SBATCH --error=logs/%x_%j.err # Error file: logs/jobname_jobid.err
#SBATCH --partition=informatik-mind # Partition to submit to (e.g., gpu, long, short)
#SBATCH --gres=gpu:1 # Request 1 GPU
#SBATCH --mem=56G # Memory per node
#SBATCH --time=168:00:00 # Time limit (hh:mm:ss)
#SBATCH --mail-type=END,FAIL # Notifications

# Module Loading
module purge
module load anaconda3/latest
. $ANACONDA_HOME/etc/profile.d/conda.sh
conda activate jax-gpu
module load cudnn/9.10-cuda12

unset DISPLAY
export MUJOCO_GL=egl
export PYOPENGL_PLATFORM=egl

xvfb-run -a python example/benchmark1.py
```