diff options
| author | Sergei Nikolaev <snikolaev@nvidia.com> | 2018-12-07 13:52:56 -0800 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2018-12-07 14:01:11 -0800 |
| commit | a0ee3a279c3eb4b8c4ef4b2553bd54ee36ecd89e (patch) | |
| tree | 1ee725a2407ba4e643e5913e018f053510dce665 /.jenkins/caffe2 | |
| parent | febc7ff99fe9759976c5357b3b67b0fefe83db17 (diff) | |
| download | pytorch-a0ee3a279c3eb4b8c4ef4b2553bd54ee36ecd89e.tar.gz pytorch-a0ee3a279c3eb4b8c4ef4b2553bd54ee36ecd89e.tar.bz2 pytorch-a0ee3a279c3eb4b8c4ef4b2553bd54ee36ecd89e.zip | |
USE_TENSORRT support and TensorRT 5 compatibility
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/13945
Differential Revision: D13317525
Pulled By: yinghai
fbshipit-source-id: 8630dfec1bbc5aac19539e344e7c38a7fd8b051d
Diffstat (limited to '.jenkins/caffe2')
| -rwxr-xr-x | .jenkins/caffe2/build.sh | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/.jenkins/caffe2/build.sh b/.jenkins/caffe2/build.sh index 89e8e39057..50e73b1e0d 100755 --- a/.jenkins/caffe2/build.sh +++ b/.jenkins/caffe2/build.sh @@ -133,6 +133,22 @@ CMAKE_ARGS+=("-DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}") if [[ $BUILD_ENVIRONMENT == *mkl* ]]; then CMAKE_ARGS+=("-DBLAS=MKL") fi + +if [[ $BUILD_ENVIRONMENT == py2-cuda9.0-cudnn7-ubuntu16.04 ]]; then + + # removing http:// duplicate in favor of nvidia-ml.list + # which is https:// version of the same repo + sudo rm -f /etc/apt/sources.list.d/nvidia-machine-learning.list + curl -o ./nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb + sudo dpkg -i ./nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb + sudo apt-key add /var/nvinfer-runtime-trt-repo-5.0.2-ga-cuda9.0/7fa2af80.pub + sudo apt-get -qq update + sudo apt-get install libnvinfer5 libnvinfer-dev + rm ./nvinfer-runtime-trt-repo-ubuntu1604-5.0.2-ga-cuda9.0_1-1_amd64.deb + + CMAKE_ARGS+=("-DUSE_TENSORRT=ON") +fi + if [[ $BUILD_ENVIRONMENT == *cuda* ]]; then CMAKE_ARGS+=("-DUSE_CUDA=ON") CMAKE_ARGS+=("-DCUDA_ARCH_NAME=Maxwell") |