Merge pull request #1067 from seemethere/torch_version

setup: Add version for torch dependency
pytorch · Nov 9, 2020 · 9ebfb7b · 9ebfb7b
2 parents 7293d43 + 1c181eb
commit 9ebfb7b
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -49,6 +49,12 @@ def _export_version(version, sha):
 
 print('-- Building version ' + VERSION)
 
+pytorch_package_version = os.getenv('PYTORCH_VERSION')
+
+pytorch_package_dep = 'torch'
+if pytorch_package_version is not None:
+    pytorch_package_dep += "==" + pytorch_package_version
+
 
 class clean(distutils.command.clean.clean):
     def run(self):
@@ -82,7 +88,7 @@ def run(self):
     license='BSD',
 
     install_requires=[
-        'tqdm', 'requests', 'torch', 'numpy'
+        'tqdm', 'requests', pytorch_package_dep, 'numpy'
     ],
     python_requires='>=3.5',
     classifiers=[