Merge pull request #149 from shenoynikhil/chunk-fix

Author: Saro00

src/boltz/model/layers/outer_product_mean.py

@@ -83,6 +83,8 @@ def forward(self, m: Tensor, mask: Tensor, chunk_size: int = None) -> Tensor:
                     z_out = z.to(m) @ sliced_weight_proj_o.T
                 else:
                     z_out = z_out + z.to(m) @ sliced_weight_proj_o.T
+            
+            z_out = z_out + self.proj_o.bias # add bias
             return z_out
         else:
             mask = mask[:, :, None, :] * mask[:, :, :, None]

tests/model/layers/test_outer_product_mean.py

@@ -21,6 +21,8 @@ def setUp(self):
         for name, param in self.layer.named_parameters():
             nn.init.normal_(param, mean=1., std=1.)
 
+        # Set to eval mode
+        self.layer.eval()
 
     def test_chunk(self):
         chunk_sizes = [16, 33, 64, 83, 100]