JAX库cuDNN初始化失败报错求助

大家好，我最近在开发一个基于JAX库的代码，不管怎么调整环境配置，都反复碰到同一个cuDNN初始化失败的错误，折腾了好久都没解决，来这儿求助各位大佬！

我的报错信息如下：

2024-08-20 16:26:58.037892: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:439] 
Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2024-08-20 16:26:58.037952: E external/xla/xla/stream_executor/cuda/cuda_dnn.cc:443] Memory usage: 46637514752 bytes free, 47587131392 bytes total.
Traceback (most recent call last):
  File "GPU_pairwise_pipline.py", line 260, in <module>
    SSMD_res_with_indices = process_blocks(train_set_sick, train_set_healthy, block_size)
  File "GPU_pairwise_pipline.py", line 172, in process_blocks
    mean_block1_sick, var_block1_sick = cal_mean_and_var(block1_sick)
  File "GPU_pairwise_pipline.py", line 17, in cal_mean_and_var
    data_jax = jnp.array(data)
  File "/a/home/cc/chemist/mikabell/anaconda3/envs/jax-env/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 2011, in array
    return _array_from_py(py_array, dtype, copy, order)
  File "/a/home/cc/chemist/mikabell/anaconda3/envs/jax-env/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1748, in _array_from_py
    aval = _abstractify(py_array)
  File "/a/home/cc/chemist/mikabell/anaconda3/envs/jax-env/lib/python3.8/site-packages/jax/_src/core.py", line 1045, in _abstractify
    return concrete_aval(x)
  File "/a/home/cc/chemist/mikabell/anaconda3/envs/jax-env/lib/python3.8/site-packages/jax/_src/core.py", line 976, in concrete_aval
    return ShapedArray(np.shape(x), dtype=np.result_type(x))
  File "/a/home/cc/chemist/mikabell/anaconda3/envs/jax-env/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 490, in __array__
    return np.asarray(self._value, dtype=dtype)

从报错里能看到，GPU明明还有大概46GB的空闲显存，但就是创建不了cuDNN句柄。我已经试过重新安装JAX、切换conda环境，但都没用。有没有大佬遇到过类似问题，能给点排查方向或者解决办法吗？

备注：内容来源于stack exchange，提问作者Mika Bell