Are there some GPU ready configurations and related documentation?

[mhmaqbool]

Hi,

I want to run the models on GPU. I was wondering if there is some GPU ready configurations or any related documentation provided for that matter?

Thanks.

[entrpn]

I haven't done extensive testing with GPUs yet. You can try running the same installation instructions, but installing jax[cuda] pip install -U "jax[cuda12]" instead of jax[tpu].

[entrpn]

@mhmaqbool did this help?