Update environment.yml (#79)

* Update environment.yml

torch==1.11.1+cu113 is not supported in pip channel now

* update other version

* add torch version check

if torch version > 1.11.1, turn on tf32

environment.yml

@@ -8,7 +8,6 @@ dependencies:
       - biopython==1.79
       - dm-tree==0.1.6
       - ml-collections==0.1.0
-      - numpy==1.21.2
       - PyYAML==5.4.1
       - requests==2.26.0
       - scipy==1.7.1
@@ -18,10 +17,12 @@ dependencies:
       - ray==2.0.0
       - pyarrow
       - pandas
-      - --find-links https://release.colossalai.org colossalai
-      - --find-links https://download.pytorch.org/whl/cu113/torch_stable.html torch==1.11.1+cu113
-      - --find-links https://download.pytorch.org/whl/cu113/torch_stable.html torchaudio==0.11.1+cu113
-      - --find-links https://download.pytorch.org/whl/cu113/torch_stable.html torchvision==0.13.1
+      - --find-links https://release.colossalai.org
+      - colossalai==0.1.10+torch1.12cu11.3
+  - pytorch::pytorch=1.12
+  - pytorch::torchvision
+  - pytorch::torchaudio
+  - conda-forge::cudatoolkit=11.3
   - conda-forge::python=3.8
   - conda-forge::setuptools=59.5.0
   - conda-forge::pip

inference.py

@@ -43,6 +43,9 @@ from fastfold.data.parsers import parse_fasta
 from fastfold.utils.import_weights import import_jax_weights_
 from fastfold.utils.tensor_utils import tensor_tree_map
 
+if int(torch.__version__.split(".")[0]) >= 1 and int(torch.__version__.split(".")[1]) > 11:
+    torch.backends.cuda.matmul.allow_tf32 = True
+
 @contextlib.contextmanager
 def temp_fasta_file(fasta_str: str):
     with tempfile.NamedTemporaryFile('w', suffix='.fasta') as fasta_file: