mypapit gnu/linux

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:

1. Docker (if you haven’t installed already)
2. NVIDIA GPU drivers for Linux (only for Ubuntu)
3. NVIDIA Container Toolkit
4. NVIDIA Tensorflow Docker containers

Docker

apt -y install docker

Installing 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

spyder

Hopefully 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:

Install Tensorflow with RTX 50 series GPU acceleration with Python wheel for Ubuntu 24.04 (and Windows WSL2)