Merge pull request #524 from shenoynikhil/attn_bug_fix

Bug fix in bias repeat when multiplicity > 1

Author: gcorso

src/boltz/model/modules/transformers.py

@@ -299,6 +299,7 @@ def forward(
             c = c.view((B * NW, W, -1))
             if mask is not None:
                 mask = mask.view(B * NW, W)
+            p = p.repeat_interleave(multiplicity, 0)
             p = p.view((p.shape[0] * NW, W, H, -1))
 
             to_keys_new = lambda x: to_keys(x.view(B, NW * W, -1)).view(B * NW, H, -1)
@@ -311,7 +312,7 @@ def forward(
             s=c,
             z=p,
             mask=mask.float(),
-            multiplicity=multiplicity,
+            multiplicity=1, # bias term already expanded with multiplicity
             to_keys=to_keys_new,
             model_cache=model_cache,
         )

src/boltz/model/modules/transformersv2.py

@@ -243,6 +243,7 @@ def forward(
         q = q.view((B * NW, W, -1))
         c = c.view((B * NW, W, -1))
         mask = mask.view(B * NW, W)
+        bias = bias.repeat_interleave(multiplicity, 0)
         bias = bias.view((bias.shape[0] * NW, W, H, -1))
 
         to_keys_new = lambda x: to_keys(x.view(B, NW * W, -1)).view(B * NW, H, -1)
@@ -253,7 +254,7 @@ def forward(
             s=c,
             bias=bias,
             mask=mask.float(),
-            multiplicity=multiplicity,
+            multiplicity=1, # bias term already expanded with multiplicity
             to_keys=to_keys_new,
         )