CUDA enabled Jupyter Docker Images

I am happy to announce that Jupyter Docker Stacks project now provides GPU accelerated Docker images. This allows PyTorch or TensorFlow operations to use compatible NVIDIA GPUs for accelerated computation.How to find these imagesAll our images are available on Quay.io registry. We provide CUDA accelerated versions of images are available for tensorflow-notebook and pytorch-notebook.To use such an image, you have to specify a special prefix tag to the image: versioned CUDA prefix like cuda11- or cuda12- for pytorch-notebook or just cuda- for tensorflow-notebook.We build pytorch-notebook only for 2 last major versions of CUDA, tensorflow-notebook image supports only the latest CUDA version listed in the officially tested build configurations list. CUDA-enabled images are available on x86_64 platform.For example, you can use an image quay.io/jupyter/pytorch-notebook:cuda12-python-3.11.8 or quay.io/jupyter/tensorflow-notebook:cuda-latest.Running the imagesTo enable CUDA support on Linux host machine, you need to:have compatible Nvidia GPUhave NVIDIA Linux driver installedadd --gpus all (or --gpus '"device=all"') flag to if you’re using Dockeradd --device 'nvidia.com/gpu=all' flag if you’re using PodmanFor example, you can run the image like this: docker run -it --rm -p 8888:8888 --gpus all quay.io/jupyter/tensorflow-notebook:cuda-latest.You can also enable GPU support on Windows using Docker or Podman.Reaching outYou will find our documentation here.If you encounter any problems, please, take a look at GitHub issues.If you would like to improve GPU support in our images, please, reach out here.AcknowledgmentsI would like to thank:Johanna Reiml, for implementing general variant concept for our images and making it work for pytorch-notebook (PR #2091)Christof Kaufmann, for making it work for tensorflow-notebook (PR #2100)Olivier Benz, for helping me out to write this blog postCUDA enabled Jupyter Docker Images was originally published in Jupyter Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.