Merge pull request #18 from gahdritz/main

Slightly improve installation process
AnnaTrainingG · Jul 3, 2022 · 01947bc · 01947bc
2 parents f66603c + af4a9ce
commit 01947bc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/flash_attn/__init__.py b/flash_attn/__init__.py
diff --git a/setup.py b/setup.py
@@ -168,5 +168,9 @@ def append_nvcc_threads(nvcc_extra_args):
     ],
     ext_modules=ext_modules,
     cmdclass={"build_ext": BuildExtension} if ext_modules else {},
-    python_requires=">=3.7"
+    python_requires=">=3.7",
+    install_requires=[
+        "torch",
+        "einops",
+    ],
 )