update by lilac

1 files changed, 33 insertions, 10 deletions

diff --git a/PKGBUILD b/PKGBUILD
index 7aa2040a0687..e8fd1283e638 100644
--- a/PKGBUILD
+++ b/PKGBUILD
@@ -4,10 +4,10 @@
 # Based on python-torchvision-git; original contributors:
 # Contributor: Stephen Zhang <zsrkmyn at gmail dot com>
 
-pkgname=python-torchvision
+pkgname=('python-torchvision' 'python-torchvision-cuda')
 _pkgname=vision
-pkgver=0.4.2
-pkgrel=7
+pkgver=0.5.0
+pkgrel=2
 pkgdesc='Datasets, transforms, and models specific to computer vision'
 arch=('x86_64')
 url='https://github.com/pytorch/vision'
@@ -37,26 +37,49 @@ checkdepends=(
   'python-scipy'
 )
 source=("${_pkgname}-${pkgver}.tar.gz::https://github.com/pytorch/vision/archive/v${pkgver}.tar.gz")
-sha512sums=('3327364bd1c10f4f0fdf9f8c1bed3848cbd887ee335634a0da6eb7a2611037408eee6261cd178d12225be5f2c809a05ff9fb87420b8accd447a799ca83def3ad')
+sha512sums=('2b76f7a0f20b7225f6310be358bf4c6fb8cdbc1fa324f53dac3bf6518276755861bd09a91da6cf6a0ff1b466104826b4094e82c796c83b9b4963ca3c0b94d93a')
 
 get_pyver() {
   python -c 'import sys; print(str(sys.version_info[0]) + "." + str(sys.version_info[1]))'
 }
 
+prepare() {
+  cp -a "${srcdir}/${_pkgname}-${pkgver}" "${srcdir}/${_pkgname}-cuda-${pkgver}"
+}
+
 build() {
-  export FORCE_CUDA=1
-  export TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;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"
-  cd "${_pkgname}-${pkgver}"
+  cd "${srcdir}/${_pkgname}-${pkgver}"
+  python setup.py build
+  
+  cd "${srcdir}/${_pkgname}-cuda-${pkgver}"
+  FORCE_CUDA=1 \
+  TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;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" \
   python setup.py build
 }
 
 check() {
-  cd "${_pkgname}-${pkgver}"
+  cd "${srcdir}/${_pkgname}-${pkgver}"
+  PYTHONPATH="${PWD}/build/lib.linux-${CARCH}-$(get_pyver)" pytest -v
+  
+  cd "${srcdir}/${_pkgname}-cuda-${pkgver}"
   PYTHONPATH="${PWD}/build/lib.linux-${CARCH}-$(get_pyver)" pytest -v
 }
 
-package() {
-  cd "${_pkgname}-${pkgver}"
+package_python-torchvision() {
+  cd "${srcdir}/${_pkgname}-${pkgver}"
+  python setup.py install --root="${pkgdir}" --optimize=1 --skip-build
+  install -Dm644 LICENSE -t "${pkgdir}/usr/share/licenses/${pkgname}"
+}
+
+package_python-torchvision-cuda() {
+  pkgdesc='Datasets, transforms, and models specific to computer vision (with GPU support)'
+  depends+=(cuda)
+  provides+=(python-torchvision=${pkgver})
+  conflicts+=(python-torchvision=${pkgver})
+  
+  cd "${srcdir}/${_pkgname}-cuda-${pkgver}"
+  FORCE_CUDA=1 \
+  TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;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" \
   python setup.py install --root="${pkgdir}" --optimize=1 --skip-build
   install -Dm644 LICENSE -t "${pkgdir}/usr/share/licenses/${pkgname}"
 }