Skip to content

GitLab

"...git@developer.sourcefind.cn:jerrrrry/infinicore.git" did not exist on "3cbe5e02150237690ff5a765e5e79376552a4c85"
Commit 29864369 authored by Gustaf Ahdritz's avatar Gustaf Ahdritz
Browse files

Squeeze templates

parent 913903e0
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
openfold/model/model.py
View file @ 29864369
......@@ -138,7 +138,7 @@ class AlphaFold(nn.Module):
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,
)
......
openfold/model/template.py
View file @ 29864369
......@@ -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,
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment