Additional fix for multimer deepspeed test

4f2f069b · Christina Floristean · Jennifer Wei · 7de0ab00 · 4f2f069b
Commit 4f2f069b authored Jan 23, 2024 by Christina Floristean Committed by Jennifer Wei May 06, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 0 additions and 1 deletion

tests/test_deepspeed_evo_attention.py tests/test_deepspeed_evo_attention.py +0 -1

No files found.
--- a/tests/test_deepspeed_evo_attention.py
+++ b/tests/test_deepspeed_evo_attention.py
@@ -306,7 +306,6 @@ class TestDeepSpeedKernel(unittest.TestCase):
        batch["residx_atom37_to_atom14"] = batch[
            "residx_atom37_to_atom14"
        ].long()
-        # print(batch["target_feat"].shape)
        batch["target_feat"] = torch.nn.functional.one_hot(batch["aatype"], consts.msa_logits - 1).to(torch.float32)
        batch["template_all_atom_mask"] = batch["template_all_atom_masks"]
        batch.update(