How to get Nvidia A100 running with TF v2

General Discussion
, ,
#1

Hello,
I have new server (Debian) with a NVIDIA A100, and cannot get tensorflow v2 (Keras) running.

python3 -m venv .venv
source .venv/bin/activate
pip install tensorflow-gpu
python run-tf.py

The same script for PyTorch works but PyTorch provides special pypi wheels that supports the Nvida A100 card. I guess something similar exists for TensorFlow but I don’t know where.

#3

I have to know How do I run TensorFlow on Nvidia GPU? Why is TensorFlow not using GPU?

#4

You don’t need -gpu. pip install tensorflow is both CPU and GPU (see TensorFlow). But I would use the docker image. Tensorflow versions 2.5 to 2.9 (current) use cuDNN v8.1 and CUDA v11.2 but that can change in the future. Being able to change TF versions without having to fidget with CUDA is great. And it makes the difference between remote and local development smaller so there’s less overhead when you switch between the two.

But what are the Python, cuDNN, and CUDA versions in your environment?