support for arm64 in pytorch-mnist image

kubeflow · Aug 29, 2022 · fced2a3 · fced2a3
1 parent 3906f8e
commit fced2a3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu
@@ -1,4 +1,6 @@
-FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
+# Pytorch=1.11.0, cuda=11.6.0
+# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel_22-08.html#rel_22-08
+FROM nvcr.io/nvidia/pytorch:22.02-py3
 
 ADD examples/v1beta1/trial-images/pytorch-mnist /opt/pytorch-mnist
 

diff --git a/scripts/v1beta1/build.sh b/scripts/v1beta1/build.sh
@@ -120,12 +120,6 @@ echo -e "\nBuilding training container images..."
 if [ ! "$ARCH" = "amd64" ]; then
   echo -e "\nSome training container images are supported only amd64."
 else
-  echo -e "\nBuilding PyTorch mnist training container example with CPU support...\n"
-  docker buildx build --platform linux/amd64 -t "${REGISTRY}/pytorch-mnist-cpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.cpu .
-
-  echo -e "\nBuilding PyTorch mnist training container example with GPU support...\n"
-  docker buildx build --platform linux/amd64 -t "${REGISTRY}/pytorch-mnist-gpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.gpu .
-
   echo -e "\nBuilding Keras CIFAR-10 CNN training container example for ENAS with GPU support...\n"
   docker buildx build --platform linux/amd64 -t "${REGISTRY}/enas-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/enas-cnn-cifar10/Dockerfile.gpu .
 
@@ -142,6 +136,12 @@ fi
 echo -e "\nBuilding mxnet mnist training container example...\n"
 docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/mxnet-mnist:${TAG}" -f examples/${VERSION}/trial-images/mxnet-mnist/Dockerfile .
 
+echo -e "\nBuilding PyTorch mnist training container example with CPU support...\n"
+docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/pytorch-mnist-cpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.cpu .
+
+echo -e "\nBuilding PyTorch mnist training container example with GPU support...\n"
+docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/pytorch-mnist-gpu:${TAG}" -f examples/${VERSION}/trial-images/pytorch-mnist/Dockerfile.gpu .
+
 echo -e "\nBuilding Tensorflow with summaries mnist training container example...\n"
 docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/tf-mnist-with-summaries:${TAG}" -f examples/${VERSION}/trial-images/tf-mnist-with-summaries/Dockerfile .