Tensorflow GPU docker image

Hello Jupyter team,

I’ve done some analysis on the docker images provided on https://github.com/jupyter/docker-stacks, and tested out different types of images for different analyses, mainly focused on the datascience and tensorflow images. However, the current tensorflow image is setup only for CPU environment. Do you have some preferred way of setting up such an image on a GPU environment, or maybe a supported image with such a setup?

I would set up this environment using one of the AWS p instances, and just install nvidia drivers, runtime and docker and start up the jupyter image there. The current jupyter/tensorflow-notebook does not have GPU support, since there is no cuda on the image. Using a docker allows for easier setup because I change instances often.

The only currently viable approach that I wanted to take is to change the most base image from ubuntu to cuda, and then follow the Dockerfile commands all the way to tensorflow, and add tensorflow-gpu to the mix. I wanted to check if there are any other options I should consider as well or if there is already a ready made solution to this problem that I haven’t had the chance to run across so far.