JAX has no gpu version

[ItsBean]

Checklist

• I have installed dependencies via poetry install (see CleanRL's installation guideline.
• [ yes] I have checked that there is no similar issue in the repo.
• [ yes] I have checked the documentation site and found not relevant information in GitHub issues.

Reviewer xxxx

In the requirements-jax.txt file, the JAX package information is listed as follows:

jax-jumpy==1.0.0 ; python_version >= "3.8" and python_version < "3.11"
jax==0.4.8 ; python_version >= "3.8" and python_version < "3.11"
jaxlib==0.4.7 ; python_version >= "3.8" and python_version < "3.11"

This configuration is for CPU users of JAX. Could the authors please provide a GPU version of the td3_jax code along with a requirements-jax-gpu.txt file? Your contribution would be greatly appreciated.

Additionally, I am unable to directly upgrade JAX to the GPU version because the td3_jax code relies heavily on Flax and Optax. Any guidance or support on this matter would be extremely helpful.

[kinalmehta]

You may refer jax documentation to install GPU version