distogram bug fix for bs > 1

Cloud303 · Aug 14, 2024 · a67a67a · a67a67a
1 parent d87a740
commit a67a67a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/models/components/attention_pair_bias.py b/src/models/components/attention_pair_bias.py
@@ -60,7 +60,7 @@ def __init__(
 
         # Attention
         self.attention = Attention(
-            c_q=dim,
+            c_q=dim,  # TODO: this Q needs to be projected from a linear layer with bias
             c_k=dim,
             c_v=dim,
             c_hidden=dim // no_heads,

diff --git a/src/utils/loss.py b/src/utils/loss.py
@@ -180,8 +180,8 @@ def distogram_loss(
     #         (eps + torch.sum(square_mask, dim=(-1, -2))))
     denom = eps + torch.sum(square_mask, dim=(-1, -2))
     mean = errors * square_mask
-    mean = mean / denom[..., None]
     mean = torch.sum(mean, dim=-1)
+    mean = mean / denom[..., None]
 
     # Average over the batch dimensions
     mean = torch.mean(mean)