Closed La-SilverLand closed 3 years ago
I got it fixed, the jax>=0.2.4 does not mean jaxlib>=0.2.4 for my situation, i have cuda 10.2 and jax==0.2.13 jaxlib==0.1.65+cuda102 is suitable
@La-SilverLand Hello! Sorry to bother you. Have you ever encountered the following problem? I think it may be related to the version of flax.
Traceback (most recent call last):
File "lra_benchmarks/listops/train.py", line 28, in <module>
from flax.deprecated import nn
ModuleNotFoundError: No module named 'flax.deprecated'
@La-SilverLand This is my setting:
in the requirement.txt it requires jax>=0.2.4 when incurring this problem as in the title, the jax github homepages says for both support of GPU and CPU, one needs to install as
but that jaxlib repo link does not contain any version >= 0.2.4, actually the newest version there is 0.1.67
How do you install the environment ?