Adds mpi4py to pip install requirements

aqlaboratory · May 6, 2024 · d75f8ae · d75f8ae
1 parent cb612f3
commit d75f8ae
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/environment.yml b/environment.yml
@@ -33,6 +33,7 @@ dependencies:
   - pytorch::pytorch=2.1
   - pytorch::pytorch-cuda=12.1
   - pip:
+      - mpi4py==3.1.5
       - deepspeed==0.12.4
       - dm-tree==0.1.6
       - git+https://github.com/NVIDIA/dllogger.git

diff --git a/tests/test_permutation.py b/tests/test_permutation.py
@@ -113,7 +113,6 @@ def test_2_permutation_pentamer(self):
 
         aligns, _ = compute_permutation_alignment(out, batch,
                                                   batch)
-        print(f"##### aligns is {aligns}")
         possible_outcome = [[(0, 1), (1, 0), (2, 3), (3, 4), (4, 2)], [(0, 0), (1, 1), (2, 3), (3, 4), (4, 2)]]
         wrong_outcome = [[(0, 1), (1, 0), (2, 4), (3, 2), (4, 3)], [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]]
         self.assertIn(aligns, possible_outcome)
@@ -163,7 +162,6 @@ def test_3_merge_labels(self):
         aligns, per_asym_residue_index = compute_permutation_alignment(out,
                                                                        batch,
                                                                        batch)
-        print(f"##### aligns is {aligns}")
         labels = split_ground_truth_labels(batch)
 
         labels = merge_labels(per_asym_residue_index, labels, aligns,