search

google-research / long-range-arena

Long Range Arena for Benchmarking Efficient Transformers
Apache License 2.0
710 stars 77 forks source link

Request about cuda version when using GPUs #43

Closed wuhaixu2016 closed 2 years ago

wuhaixu2016 commented 2 years ago

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 commented 2 years ago

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

BalloutAI commented 2 years ago

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

DaShenZi721 commented 1 year ago

@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 commented 1 year ago

@wuhaixu2016 This is my setting: