fix deepspeed_evo_attention to work in both monomer and multimer settings.

tests/test_deepspeed_evo_attention.py

@@ -244,9 +244,6 @@ class TestDeepSpeedKernel(unittest.TestCase):
         pair_act = np.random.rand(n_res, n_res, consts.c_z).astype(np.float32)
         pair_mask = np.random.randint(0, 2, (n_res, n_res)).astype(np.float32)
 
-        inds = np.random.randint(0, 21, (n_res,))
-        batch["target_feat"] = np.eye(22)[inds]
-
         batch = {k: torch.as_tensor(v).cuda() for k, v in batch.items()}
         template_feats = {
             k: v for k, v in batch.items() if k.startswith("template_")
@@ -309,7 +306,8 @@ class TestDeepSpeedKernel(unittest.TestCase):
         batch["residx_atom37_to_atom14"] = batch[
             "residx_atom37_to_atom14"
         ].long()
-        batch["target_feat"] = torch.nn.functional.one_hot(batch["aatype"], 21).to(torch.float32)
+        # 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(
             data_transforms.atom37_to_torsion_angles("template_")(batch)