diff --git a/py/build_whl.sh b/py/build_whl.sh
index 7706994727..9de39509d1 100755
--- a/py/build_whl.sh
+++ b/py/build_whl.sh
@@ -5,6 +5,7 @@
 cd /workspace/Torch-TensorRT/py
 
 export CXX=g++
+export CUDA_HOME=/usr/local/cuda-11.3
 
 build_py36() {
     /opt/python/cp36-cp36m/bin/python -m pip install -r requirements.txt
@@ -33,4 +34,4 @@ build_py39() {
 build_py36
 build_py37
 build_py38
-build_py39
\ No newline at end of file
+build_py39