Hi,can I run berax on my GPU ?

[NickShanyt]

Hi, I have successfully run BERTax on the CPU, but at a very limited speed.
I now have a GPU server with CUDA version 11.4.1 and CUDNN version 8.2.4, and I tried to run BERTax on it, but it failed.
After my troubleshooting, I surmise that it has something to do with the version of Tensorflow that BERTax uses.
Are you able to come up with a more detailed configuration method?

I found that CUDA version 11.4 needs to match Tensorflow version 2.6.

[NickShanyt]

Extra info, when I first install Tensorflow==2.12.0 (it works on my another CPU platform) and then install bertax via conda install -c fkretschmer bertax, bertax runs fine on the CPU.

[f-kretschmer]

Unfortunately BERTax needs an older version of tensorflow due to its keras-bert-dependency. Perhaps the docker version can help here: https://github.com/f-kretschmer/bertax#docker ?

[JiayiJennie]

Hi! Could you provide the tensorflow version used in bertax? Thanks!

[f-kretschmer]

In the development we used version 2.4.1, although everything depends on using keras-bert, with the version 0.86.0.