support arm64 in darts-cnn-cifar10 image

Signed-off-by: tenzen-y <[email protected]>

Author: tenzen-y

examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu

@@ -1,6 +1,6 @@
-ARG cuda_version=10.0
-ARG cudnn_version=7
-FROM pytorch/pytorch:1.0-cuda${cuda_version}-cudnn${cudnn_version}-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
 
 ENV TARGET_DIR /opt/darts-cnn-cifar10
 

scripts/v1beta1/build.sh

@@ -120,16 +120,16 @@ 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 CIFAR-10 CNN training container example for DARTS with CPU support...\n"
-  docker buildx build --platform linux/amd64 -t "${REGISTRY}/darts-cnn-cifar10-cpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.cpu .
-
-  echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with GPU support...\n"
-  docker buildx build --platform linux/amd64 -t "${REGISTRY}/darts-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.gpu .
-
   echo -e "\nBuilding dynamic learning rate training container example for PBT...\n"
   docker buildx build --platform linux/amd64 -t "${REGISTRY}/simple-pbt:${TAG}" -f examples/${VERSION}/trial-images/simple-pbt/Dockerfile .
 fi
 
+echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with CPU support...\n"
+docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/darts-cnn-cifar10-cpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.cpu .
+
+echo -e "\nBuilding PyTorch CIFAR-10 CNN training container example for DARTS with GPU support...\n"
+docker buildx build --platform "linux/$ARCH" -t "${REGISTRY}/darts-cnn-cifar10-gpu:${TAG}" -f examples/${VERSION}/trial-images/darts-cnn-cifar10/Dockerfile.gpu .
+
 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 .