Open astanziola opened 2 years ago
@djps FYI, not sure if you want to give it a try.
I don't have a Windows with CUDA machine at hand now, but will test this out as soon as I can
I tried this but had conflicts between the version of cuda I have and the version it wanted to install
There are community-driven wheels for
jax
andjaxlib
with GPU support viaconda-forge
Relevant links:
I managed to get a working installation of
jax
on with onlyconda
installed (no cuda or cudnn prior installation) usingIt is interesting to see if this works under windows, to avoid using the WSL or building
jax
from scratch