From 1317e855751e0a16666948068b7b31f824b1312b Mon Sep 17 00:00:00 2001
From: chand1012 <chandler@chand1012.dev>
Date: Thu, 15 Jul 2021 01:25:28 +0000
Subject: [PATCH] Update dockerfiles

---
 Dockerfile.arm64 | 5 ++---
 Dockerfile.cpu   | 2 +-
 Dockerfile.gpu   | 2 +-
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/Dockerfile.arm64 b/Dockerfile.arm64
index 432772b16..2f606f95a 100644
--- a/Dockerfile.arm64
+++ b/Dockerfile.arm64
@@ -1,4 +1,4 @@
-FROM chand1012/tensorflow:2.5.0-py3
+FROM chand1012/tensorflow:2.5.0
 
 RUN mkdir /gpt-2
 WORKDIR /gpt-2
@@ -6,5 +6,4 @@ ADD . /gpt-2
 RUN pip3 install -r requirements.txt
 RUN python3 download_model.py 124M
 RUN python3 download_model.py 355M
-RUN python3 download_model.py 774M
-RUN python3 download_model.py 1558M
+# additional models would not build on the Pi
diff --git a/Dockerfile.cpu b/Dockerfile.cpu
index 9877ece78..6a3a4a1e9 100644
--- a/Dockerfile.cpu
+++ b/Dockerfile.cpu
@@ -1,4 +1,4 @@
-FROM tensorflow/tensorflow:2.5.0-py3
+FROM tensorflow/tensorflow:2.5.0
 
 ENV LANG=C.UTF-8
 RUN mkdir /gpt-2
diff --git a/Dockerfile.gpu b/Dockerfile.gpu
index 181eaa192..143031194 100644
--- a/Dockerfile.gpu
+++ b/Dockerfile.gpu
@@ -1,4 +1,4 @@
-FROM tensorflow/tensorflow:2.5.0-gpu-py3
+FROM tensorflow/tensorflow:2.5.0-gpu
 
 # nvidia-docker 1.0
 LABEL com.nvidia.volumes.needed="nvidia_driver"