made FA a global package

Calcifer26 · Aug 1, 2024 · ab5b56b · ab5b56b
1 parent 8f5207e
commit ab5b56b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -3,10 +3,10 @@ packaging~=23.2
 absl-py
 
 # --------- pytorch --------- #
-torch~=2.2.0
-torchvision~=0.17.0
-lightning~=2.1.4
-torchmetrics~=1.3.0.post0
+torch==2.3.1
+torchvision
+lightning
+torchmetrics
 
 # --------- kernels --------- #
 deepspeed  # works in Colab

diff --git a/setup.py b/setup.py
@@ -138,7 +138,7 @@ def run_flash_attn_setup():
     author_email='[email protected]',
     license='Apache License, Version 2.0',
     url='https://github.com/Ligo-Biosciences/AlphaFold3',
-    packages=find_packages(exclude=["tests", "scripts"]),
+    packages=find_packages(include=["flash_attn", "flash_attn.*"], exclude=["tests", "scripts"]),
     include_package_data=True,
     package_data={
         "src": ['utils/kernel/csrc/*'],