From 8c3fd8bac3e5a54a20e860fc163d6e67cc03c623 Mon Sep 17 00:00:00 2001
From: Janne Hellsten <jhellsten@nvidia.com>
Date: Mon, 26 Oct 2020 14:10:54 +0200
Subject: [PATCH] Default to nvcr.io/nvidia/tensorflow:20.10-tf1-py3 base image
 (#10)

This works out of the box on newer NVIDIA GPUs such as RTX 3090.
Note that an NVIDIA driver release r455.23 or later is required to
run this image.

See README for instructions on how to build for older drivers.
---
 Dockerfile | 3 ++-
 README.md  | 6 ++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index d29740ee..51817e50 100755
--- a/Dockerfile
+++ b/Dockerfile
@@ -6,7 +6,8 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-FROM tensorflow/tensorflow:1.14.0-gpu-py3
+ARG BASE_IMAGE=nvcr.io/nvidia/tensorflow:20.10-tf1-py3
+FROM $BASE_IMAGE
 
 RUN pip install scipy==1.3.3
 RUN pip install requests==2.22.0
diff --git a/README.md b/README.md
index 9b3d0d25..bbe52379 100755
--- a/README.md
+++ b/README.md
@@ -88,6 +88,12 @@ docker run --gpus all -it --rm -v `pwd`:/scratch --user $(id -u):$(id -g) styleg
     --outdir=out --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl)"
 ```
 
+**Note**: The above defaults to a container base image that requires NVIDIA driver release `r455.23` or later.  To build an image for older drivers and GPUs, run:
+
+```.bash
+docker build --build-arg BASE_IMAGE=tensorflow/tensorflow:1.14.0-gpu-py3 --tag stylegan2ada:latest .
+```
+
 ## Projecting images to latent space
 
 To find the matching latent vector for a given image file, run: