fix(//py): Constrain the CUDA version in container builds

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Nov 10, 2021 · a21a045 · a21a045
1 parent b8bf66e
commit a21a045
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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
+build_py39