From d75f8ae402f0de3c9de9015efbe823374d6d4d28 Mon Sep 17 00:00:00 2001 From: jnwei Date: Mon, 6 May 2024 10:48:06 +0700 Subject: [PATCH] Adds mpi4py to pip install requirements --- environment.yml | 1 + tests/test_permutation.py | 2 -- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/environment.yml b/environment.yml index fa4a4e634..c15aa2401 100644 --- 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 index d0db977ae..740de445a 100644 --- 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,