Adds CUDA 10.0 support to pip package

apache · Feb 25, 2019 · f0713ce · f0713ce
1 parent e2de496
commit f0713ce
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tools/pip/setup.py b/tools/pip/setup.py
@@ -112,7 +112,9 @@ def has_ext_modules(self):
 if variant == 'CPU':
     libraries.append('openblas')
 else:
-    if variant.startswith('CU92'):
+    if variant.startswith('CU100'):
+        libraries.append('CUDA-10.0')
+    elif variant.startswith('CU92'):
         libraries.append('CUDA-9.2')
     elif variant.startswith('CU91'):
         libraries.append('CUDA-9.1')