It seems that the Tensorflow acceleration is broken with the latest RTX50 series GPU, especially with the most cost effective RTX5060Ti 16GB card.
This is a how to guide to re-enable the Tensorflow acceleration with the official Tensorflow docker image from NVIDIA
Basically you need to install these four things:
- Docker (if you haven’t installed already)
- NVIDIA GPU drivers for Linux (only for Ubuntu)
- NVIDIA Container Toolkit
- NVIDIA Tensorflow Docker containers
Docker
apt -y install dockerInstalling NVIDIA GPU Driver
After downloading the NVIDIA GPU drivers for Linux (only for fully Linux-based operating system). This step is unnecessary for Microsoft Windows WSL2
bash NVIDIA_Linux_x86_64 570.153.02.run
Installing NVIDIA Container Toolkit
Follow the instruction here: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html
Do not forget to configure docker
sudo nvidia-ctk runtime configure --runtime=docker
Then restart docker to enable gpu driver integration
sudo systemctl restart docker
Test the nvidia driver under the container runtime
sudo docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi
Install NVIDIA Tensorflow docker containers (basic)
docker run --gpus all -it --rm nvcr.io/nvidia/tensorflow:xx.xx-tfx-py3
Replace the xx with the actualy version of the tensorflow containers, at the time of the writing it is:
docker run --gpus all -it --rm nvcr.io/nvidia/tensorflow:25.02-tf2-py3
If you want to run it inside the docker, and link your /home directory with the /workspace directory inside the docker image, you can just run :
docker run --gpus all -it --rm -v /home/username:/workspace nvcr.io/nvidia/tensorflow:25.02-tf2-py3
BONUS: How to install Spyder, Jupyterlab and additional Tensorflow/Keras libraries in NVIDIA docker image
I’ve prepared a Dockerfile to rebuild the NVIDIA Tensorflow docker container with GPU acceleration. Download the dockerfile and run “docker build” with this parameter
docker build -t my-nvidia-tf-ds .
Then you can run the container with GPU acceleration
docker run --gpus all -it --rm -v /root:/workspace my-nvidia-tf-ds
You can also expose the port and run Jupyterlab within the docker image. Just follow this step:
docker run --gpus all -p 8888:8888 -it my-nvidia-tf-ds \
jupyter lab --ip=0.0.0.0 --allow-root
Additionally, you can run Spyder in the docker image by forwarding xhost :
xhost +local:docker
docker run -it --gpus all -e DISPLAY=$DISPLAY -v /tmp/.X11-unix:/tmp/.X11-unix \
-v /root:/workspace my-nvidia-tf-ds
spyderHopefully this would help you run GPU accelerated Tensorflow with RTX50 series GPU card.
This also works under Microsoft Windows 11 / WSL2 environment too!
If you want to install Tensorflow with RTX 50 series support directly inside your Ubuntu environment, then refer to this post: