Request about cuda version when using GPUs

[wuhaixu2016]

Hi, Thanks for your code.

I try to run this code on GPUs. But I do suffer a lot in the environment setting.

Can you provide the version of CUDA, jaxlib, jax, flax, when using GPUs for training?

[wuhaixu2016]

Hi, I have solved this problem with the following settings:

• cuda: 10.1
• jax: 0.2.13
• jaxlib: 0.1.65+cuda101
• flax: 0.2.2
• tensorflow: 2.7.0

[BalloutAI]

Hello,
How were you able to install cuda 10.1 with tensorflow 2.7.0 as it needs cuda 11 version

[DaShenZi721]

@wuhaixu2016 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'

[DaShenZi721]

@wuhaixu2016 This is my setting:

• cuda: 11.2
• jax: 0.2.13
• jaxlib: 0.1.65+cuda112
• flax: 0.2.2
• tensorflow: 2.7.0