Jax version does not find GPU about moscot HOT 8 CLOSED

on hold.

from moscot.

this is still a problem for me...

from moscot.

I solved it with this

pip install --upgrade pip
# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html  # Note: wheels only available on linux.

which is btw installation recommended in jax readme. Can we add it to requirements?

from moscot.

@MUCDK, could you take care of this please?

from moscot.

Done. Reopen if you have problems

from moscot.

The internet says there is no easy solution to this as adding a --global option in requirements.txt file does not confluence with the setup.py file. Maybe @michalk8 knows an elegant solution?

from moscot.

The internet says there is no easy solution to this as adding a --global option in requirements.txt file does not confluence with the setup.py file. Maybe @michalk8 knows an elegant solution?

Unfortunately, nothing comes to mind; would just use CPU version in reqs and for GPU point to jax docs.

from moscot.

Think this can be closed, feel free to re-open.

from moscot.