CUDA and Jax version

Open cerabinowitz opened this issue 1 year ago • 1 comments

Hi,

Thank you so much for this repo! What cuda version were you using when you trained with GPUs and did you use a specific Jax version as well?

Thank you!

May 27 '24 16:05 cerabinowitz

Glad you like it!

Here's what we use:

- jax==0.4.26
- jaxlib==0.4.26+cuda12.cudnn89

Hope that helps?

May 30 '24 04:05 reinerp