added batch_size dimesion to compute_tm

15105078 · Geoffrey Yu · a3ea7c65 · 15105078
Commit 15105078 authored Jun 28, 2023 by Geoffrey Yu
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

openfold/utils/loss.py openfold/utils/loss.py +2 -2

No files found.
--- a/openfold/utils/loss.py
+++ b/openfold/utils/loss.py
@@ -694,7 +694,7 @@ def compute_tm(
    predicted_tm_term = torch.sum(probs * tm_per_bin, dim=-1)

    n = residue_weights.shape[-1]
-    pair_mask = residue_weights.new_ones((n, n), dtype=torch.int32)
+    pair_mask = residue_weights.new_ones((1,n, n), dtype=torch.int32)
    if interface:
        pair_mask *= (asym_id[..., None] != asym_id[..., None, :]).to(dtype=pair_mask.dtype)