Bug fixes in vqvae module

src/vqvae/attention.py

@@ -23,7 +23,7 @@ def __init__(
 
         self.d_head = self.d_model // self.n_heads
         self.layernorm_qkv = nn.Sequential(
-            nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=bias)
+            nn.LayerNorm(d_model, bias=bias), nn.Linear(d_model, d_model * 3, bias=bias)
         )
         self.out_proj = nn.Linear(d_model, d_model, bias=bias)
 

src/vqvae/quantizer_module.py

@@ -56,7 +56,7 @@ def get_codebook(self,):
             return self.codebook.weight
 
     def indices2embedding(self, indices: torch.IntTensor) -> torch.Tensor:
-        z_q = self.codebook[indices]
+        z_q = self.codebook(indices)
         return z_q
 
     def forward(self, z: torch.Tensor) -> (torch.Tensor, torch.IntTensor, float):

src/vqvae/transformer_stack.py

@@ -63,10 +63,10 @@ def __init__(
     def forward(
         self,
         x: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        affine: Affine3D,
+        affine_mask: torch.Tensor,
+        attention_mask: torch.Tensor,
         sequence_id: torch.Tensor | None = None,
-        affine: Affine3D | None = None,
-        affine_mask: torch.Tensor | None = None,
         chain_id: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
@@ -75,8 +75,8 @@ def forward(
         Args:
             x (torch.Tensor): The input tensor of shape (batch_size, sequence_length, d_model).
             sequence_id (torch.Tensor): The sequence ID tensor of shape (batch_size, sequence_length).
-            affine (Affine3D | None): The affine transformation tensor or None.
-            affine_mask (torch.Tensor | None): The affine mask tensor or None.
+            affine (Affine3D): The affine transformation tensor.
+            affine_mask (torch.Tensor): The affine mask tensor.
             chain_id (torch.Tensor): The protein chain tensor of shape (batch_size, sequence_length).
                 Only used in geometric attention.
 
@@ -89,4 +89,4 @@ def forward(
             chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device)
         for block in self.blocks:
             x = block(x, attention_mask, sequence_id, affine, affine_mask, chain_id)
-        return self.norm(x), x
+        return self.norm(x), x