Squeeze templates

aqlaboratory · Aug 19, 2022 · 2986436 · 2986436
1 parent 913903e
commit 2986436
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/openfold/model/model.py b/openfold/model/model.py
@@ -138,7 +138,7 @@ def embed_templates(self, batch, z, pair_mask, templ_dim, inplace_safe):
         for i in range(n_templ):
             idx = batch["template_aatype"].new_tensor(i)
             single_template_feats = tensor_tree_map(
-                lambda t: torch.index_select(t, templ_dim, idx),
+                lambda t: torch.index_select(t, templ_dim, idx).squeeze(templ_dim),
                 batch,
             )
 

diff --git a/openfold/model/template.py b/openfold/model/template.py
@@ -448,7 +448,7 @@ def embed_templates_offload(
     for i in range(n_templ):
         idx = batch["template_aatype"].new_tensor(i)
         single_template_feats = tensor_tree_map(
-            lambda t: torch.index_select(t, templ_dim, idx),
+            lambda t: torch.index_select(t, templ_dim, idx).squeeze(templ_dim),
             batch,
         )