Update requirements

environment.yml

@@ -4,6 +4,17 @@ channels:
   - bioconda
   - pytorch
 dependencies:
+  - conda-forge::python=3.7
+  - conda-forge::setuptools=59.5.0
+  - conda-forge::pip
+  - conda-forge::openmm=7.5.1
+  - conda-forge::pdbfixer
+  - conda-forge::cudatoolkit==10.2.*
+  - conda-forge::cudatoolkit-dev==10.*
+  - bioconda::hmmer==3.3.2
+  - bioconda::hhsuite==3.3.0
+  - bioconda::kalign2==2.04
+  - pytorch::pytorch=1.12.*
   - pip:
       - biopython==1.79
       - deepspeed==0.5.9
@@ -16,15 +27,6 @@ dependencies:
       - tqdm==4.62.2
       - typing-extensions==3.10.0.2
       - pytorch_lightning==1.5.10
+      - wandb==0.12.21
+      - git+https://github.com/HazyResearch/flash-attention.git
       - git+https://github.com/NVIDIA/dllogger.git
-  - pytorch::pytorch=1.10.*
-  - conda-forge::python=3.7
-  - conda-forge::setuptools=59.5.0
-  - conda-forge::pip
-  - conda-forge::openmm=7.5.1
-  - conda-forge::pdbfixer
-  - conda-forge::cudatoolkit==10.2.*
-  - conda-forge::cudatoolkit-dev==10.*
-  - bioconda::hmmer==3.3.2
-  - bioconda::hhsuite==3.3.0
-  - bioconda::kalign2==2.04