adds mpi4py to environment and removes stray print statements from test

environment.yml

@@ -32,6 +32,7 @@ dependencies:
   - bioconda::kalign2==2.04
   - pytorch::pytorch=1.12.*
   - pip:
+      - mpi4py==3.1.5
       - deepspeed==0.12.4
       - dm-tree==0.1.6
       - git+https://github.com/NVIDIA/dllogger.git

tests/test_permutation.py

@@ -113,7 +113,6 @@ class TestPermutation(unittest.TestCase):
 
         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 @@ class TestPermutation(unittest.TestCase):
         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,