Home Setting up PyTorch Development in Docker Compose for Ubuntu 22.04

Setting up PyTorch Development in Docker Compose for Ubuntu 22.04

There’s a fantastic docker-compose setup by cresset which takes care of most of the setup for you. However, if you encounter an error message indicating that Docker is unable to utilize the NVIDIA graphics card on your machine, it is likely due to missing NVIDIA drivers or an improperly configured nvidia-docker plugin.

To resolve this issue, follow these steps:

Install NVIDIA Drivers

First, verify if you have the appropriate NVIDIA drivers installed by executing the following command in your terminal:


If the command runs successfully, you have the necessary NVIDIA drivers installed and should output their version.

To install the NVIDIA drivers, run the following commands:

sudo ubuntu-drivers autoinstall

Install the nvidia-docker Plugin

The nvidia-docker plugin enables Docker to interact with your NVIDIA GPU. To install it, run the following commands:

curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/ubuntu22.04/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

sudo apt-get update
sudo apt-get install -y nvidia-docker2

After the installation completes, restart the Docker daemon to apply the changes:

sudo systemctl restart docker

Verify the Installation

To verify that the NVIDIA GPU is now accessible within Docker, you can run the following command:

docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

replacing 11.8.0 with the version of CUDA you want to use (obtained from nvidia-smi) output.

If everything is set up correctly, you should see the output of the nvidia-smi command within the Docker container.

For more information about available Docker images, you can visit the NVIDIA CUDA Docker Hub page.

Now you have successfully set up PyTorch development in Docker Compose with NVIDIA GPU support on Ubuntu 22.04. Happy coding!

This post is licensed under CC BY 4.0 by the author.