add more env_vars back in

2 files changed, 27 insertions, 22 deletions

diff --git a/.SRCINFO b/.SRCINFO
index 32b9f4280828..0979f63e6a7d 100644
--- a/.SRCINFO
+++ b/.SRCINFO
@@ -1,7 +1,7 @@
 pkgbase = python-pytorch-rocm
 	pkgdesc = Tensors and Dynamic neural networks in Python with strong GPU acceleration
 	pkgver = 1.6.0
-	pkgrel = 5
+	pkgrel = 6
 	url = https://pytorch.org
 	arch = x86_64
 	license = BSD
diff --git a/PKGBUILD b/PKGBUILD
index 39bf535a6bfe..081d4b38006b 100644
--- a/PKGBUILD
+++ b/PKGBUILD
@@ -7,7 +7,7 @@ pkgname=("python-pytorch-rocm" "python-pytorch-opt-rocm")
 _pkgname="pytorch"
 pkgver=1.6.0
 _pkgver=1.6.0
-pkgrel=5
+pkgrel=6
 pkgdesc="Tensors and Dynamic neural networks in Python with strong GPU acceleration"
 arch=('x86_64')
 url="https://pytorch.org"
@@ -60,36 +60,33 @@ prepare() {
   # remove local nccl
   rm -rf third_party/nccl/nccl
 
-  # Apply changes needed for ROCm
-  python tools/amd_build/build_amd.py
-
   cd ..
 
   cp -a "${_pkgname}-${pkgver}" "${_pkgname}-${pkgver}-rocm"
   cp -a "${_pkgname}-${pkgver}" "${_pkgname}-${pkgver}-opt-rocm"
 
-  #export VERBOSE=1
-  #export PYTORCH_BUILD_VERSION="${pkgver}"
-  #export PYTORCH_BUILD_NUMBER=1
+  export VERBOSE=1
+  export PYTORCH_BUILD_VERSION="${pkgver}"
+  export PYTORCH_BUILD_NUMBER=1
 
   # Check tools/setup_helpers/cmake.py, setup.py and CMakeLists.txt for a list of flags that can be set via env vars.
-  #export USE_MKLDNN=ON
-  ## export BUILD_CUSTOM_PROTOBUF=OFF
-  ## export BUILD_SHARED_LIBS=OFF
+  export USE_MKLDNN=ON
+  # export BUILD_CUSTOM_PROTOBUF=OFF
+  # export BUILD_SHARED_LIBS=OFF
   #export USE_FFMPEG=ON
-  #export USE_GFLAGS=ON
-  #export USE_GLOG=ON
+  export USE_GFLAGS=ON
+  export USE_GLOG=ON
   #export BUILD_BINARY=ON
   #export USE_OPENCV=ON
-  #export USE_SYSTEM_NCCL=ON
-  #export NCCL_VERSION=$(pkg-config nccl --modversion)
-  #export NCCL_VER_CODE=$(sed -n 's/^#define NCCL_VERSION_CODE\s*\(.*\).*/\1/p' /usr/include/nccl.h)
-  #export CUDAHOSTCXX=g++-9
-  #export CUDA_HOME=/opt/cuda
-  #export CUDNN_LIB_DIR=/usr/lib
-  #export CUDNN_INCLUDE_DIR=/usr/include
-  #export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
-  #export TORCH_CUDA_ARCH_LIST="5.2;5.3;6.0;6.0+PTX;6.1;6.1+PTX;6.2;6.2+PTX;7.0;7.0+PTX;7.2;7.2+PTX;7.5;7.5+PTX;8.0;8.0+PTX;"
+  export USE_SYSTEM_NCCL=ON
+  export NCCL_VERSION=$(pkg-config nccl --modversion)
+  export NCCL_VER_CODE=$(sed -n 's/^#define NCCL_VERSION_CODE\s*\(.*\).*/\1/p' /usr/include/nccl.h)
+  export CUDAHOSTCXX=g++-9
+  export CUDA_HOME=/opt/cuda
+  export CUDNN_LIB_DIR=/usr/lib
+  export CUDNN_INCLUDE_DIR=/usr/include
+  export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
+  export TORCH_CUDA_ARCH_LIST="5.2;5.3;6.0;6.0+PTX;6.1;6.1+PTX;6.2;6.2+PTX;7.0;7.0+PTX;7.2;7.2+PTX;7.5;7.5+PTX;8.0;8.0+PTX;"
 }
 
 build() {
@@ -97,6 +94,10 @@ build() {
   export USE_CUDA=OFF
   #export USE_ROCM=ON
   cd "${srcdir}/${_pkgname}-${pkgver}-rocm"
+
+  # Apply changes needed for ROCm
+  python tools/amd_build/build_amd.py
+
   python setup.py build
 
 
@@ -105,6 +106,10 @@ build() {
   #export USE_ROCM=ON
   cd "${srcdir}/${_pkgname}-${pkgver}-opt-rocm"
   echo "add_definitions(-march=haswell)" >> cmake/MiscCheck.cmake
+
+  # Apply changes needed for ROCm
+  python tools/amd_build/build_amd.py
+
   python setup.py build
 }