diff --git a/Dockerfile b/Dockerfile
index 3b46f20..a76ccae 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,21 +3,24 @@ FROM nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.6-py3
 ENV DEBIAN_FRONTEND=noninteractive
 
 RUN apt update
-RUN apt install -y cmake libgtk2.0-dev wget
+RUN apt install -y cmake libgtk2.0-dev wget python3.8 python3.8-dev
 # ffmpeg
 RUN apt install -y libavcodec-dev libavformat-dev libavutil-dev libswscale-dev libavresample3
 # gstreamer
 RUN apt install -y libgstreamer-opencv1.0-0 libgstreamer1.0-dev libgstreamer-plugins-base1.0-dev libgstreamer1.0-dev
 
-RUN wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
-RUN python3 -m pip install torch-1.8.0-cp36-cp36m-linux_aarch64.whl
+RUN python3.8 -m pip install --upgrade pip
+RUN python3.8 -m pip install setuptools gdown
+RUN gdown https://drive.google.com/uc?id=1V231Nmx42vXTo5nq_YsV_BouZwBE9vjh
+# RUN wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
+RUN python3.8 -m pip install torch-1.8.0a0+56b43f4-cp38-cp38-linux_aarch64.whl
 RUN git clone https://github.com/pytorch/vision torchvision
-RUN cd torchvision && git checkout v0.9.0 && python3 setup.py install
+RUN apt install libomp5
+RUN cd torchvision && git checkout v0.9.0 && python3.8 setup.py install
 
-RUN python3 -m pip install scikit-build
+RUN python3.8 -m pip install scikit-build
 RUN git clone --recursive https://github.com/skvark/opencv-python.git
-RUN python3 -m pip install --upgrade pip
-RUN cd opencv-python && python3 -m pip wheel . --verbose && find . -name "opencv_python*.whl" | xargs python3 -m pip install
+RUN cd opencv-python && python3.8 -m pip wheel . --verbose && find . -name "opencv_python*.whl" | xargs python3.8 -m pip install
 RUN git clone https://github.com/ultralytics/yolov5.git
-RUN cd yolov5 && python3 -m pip install -r requirements.txt
+RUN cd yolov5 && python3.8 -m pip install -r requirements.txt
 WORKDIR /yolov5
diff --git a/build.sh b/build.sh
index 88a0cc0..c21370b 100755
--- a/build.sh
+++ b/build.sh
@@ -1,8 +1,14 @@
 #!/bin/bash
+set -e
+
+if [ ! -f /etc/nv_tegra_release ]; then
+    echo "Error: $0 should be run on the Jetson platform."
+    exit 1
+fi
 if [[ $(cat /etc/nv_tegra_release) =~ ^.*REVISION:[^\S]([0-9]*\.[0-9]).*$ ]]; then
     case ${BASH_REMATCH[1]} in
         5.*) sudo docker build -t yolov5 . ;;
-        4.*) sudo docker build -t yolov5 -f Dockerfile.jetpack4.4 . ;;
-        * ) echo "unknown jetpack ${BASH_REMATCH[1]}" ;;
+        * ) echo "unknown jetpack ${BASH_REMATCH[1]}"
+            exit 1 ;;
     esac
 fi
diff --git a/run.sh b/run.sh
index 47b9bd2..8e1c4ab 100755
--- a/run.sh
+++ b/run.sh
@@ -1,4 +1,4 @@
 #!/bin/bash
 xhost +local:
-sudo docker run -it --rm --runtime nvidia --network host --device /dev/video0:/dev/video0:mrw -e DISPLAY=$DISPLAY -v /tmp/.X11-unix/:/tmp/.X11-unix yolov5 python3 detect.py --source 0
+sudo docker run -it --rm --runtime nvidia --network host --device /dev/video0:/dev/video0:mrw -e DISPLAY=$DISPLAY -v /tmp/.X11-unix/:/tmp/.X11-unix yolov5 python3.8 detect.py --source 0