Merge branch 'main' of github.com:jwohlwend/boltz

Author: jwohlwend

src/boltz/model/layers/outer_product_mean.py

@@ -85,6 +85,7 @@ def forward(self, m: Tensor, mask: Tensor, chunk_size: int = None) -> Tensor:
                     z_out = z_out + z.to(m) @ sliced_weight_proj_o.T
             return z_out
         else:
+            mask = mask[:, :, None, :] * mask[:, :, :, None]
             num_mask = mask.sum(1).clamp(min=1)
             z = torch.einsum("bsic,bsjd->bijcd", a.float(), b.float())
             z = z.reshape(*z.shape[:3], -1)