diff --git a/docker/build.sh b/docker/build.sh index f09767408..795613b54 100755 --- a/docker/build.sh +++ b/docker/build.sh @@ -16,7 +16,11 @@ BASE_IMAGE=$1 source tools/l4t-version.sh if [ -z $BASE_IMAGE ]; then - if [ $L4T_VERSION = "32.4.4" ]; then + if [ $L4T_VERSION = "32.5.1" ]; then + BASE_IMAGE="nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.6-py3" + elif [ $L4T_VERSION = "32.5.0" ]; then + BASE_IMAGE="nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.6-py3" + elif [ $L4T_VERSION = "32.4.4" ]; then BASE_IMAGE="nvcr.io/nvidia/l4t-pytorch:r32.4.4-pth1.6-py3" elif [ $L4T_VERSION = "32.4.3" ]; then BASE_IMAGE="nvcr.io/nvidia/l4t-pytorch:r32.4.3-pth1.6-py3"