Merge pull request #13 from eitanturok/main

Cleanup `setup.py`
stanford-futuredata · Jul 29, 2024 · a1ddf98 · a1ddf98
2 parents e5c47f6 + 31ca71b
commit a1ddf98
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/setup.py b/setup.py
@@ -1,13 +1,17 @@
-from setuptools import setup, find_packages
-import torch
+from setuptools import find_packages, setup
+
+install_requires = [
+    'torch>=2.3.0,<2.4',
+    'triton>=2.1.0',
+]
 
 extra_deps = {}
 
 extra_deps['dev'] = [
     'absl-py',
 ]
 
-extra_deps['all'] = set(dep for deps in extra_deps.values() for dep in deps)
+extra_deps['all'] = list(set(dep for deps in extra_deps.values() for dep in deps))
 
 setup(
     name="stanford-stk",
@@ -24,8 +28,6 @@
         "Operating System :: Unix",
     ],
     packages=find_packages(),
-    install_requires=[
-        "triton>=2.1.0",
-    ],
+    install_requires=install_requires,
     extras_require=extra_deps,
 )