Problems with jax #98
Description
I installed the required libraries by pip install -r requirement.txt. The CUDA works well and the GPU can be found by tensorflow. However, when I try to run the code, an error occurs.
"Unable to initialize backend 'cuda'": module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I searched on the web and found it might be caused by the version of the libraries.
Could you please share the versions of those packages you used with me? Thank you very much.
