msa

Genentech · Apr 29, 2024 · 6923048 · 6923048
1 parent 2d04fc9
commit 6923048
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/src/beignet/nn/_msa.py b/src/beignet/nn/_msa.py
@@ -1,6 +1,6 @@
 import torch
 from torch import Tensor
-from torch.nn import Conv2d, Module
+from torch.nn import Conv1d, Module
 
 import beignet.operators
 
@@ -21,12 +21,17 @@ def __init__(
 
         self.temperature = temperature
 
-        self.embed = Conv2d(in_channels, out_channels, kernel_size)
+        self.embedding = Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            padding="same",
+        )
 
     def forward(self, inputs: (Tensor, Tensor)) -> Tensor:
         matrices, shapes = inputs
 
-        embedding = self.embed(matrices)
+        embedding = self.embedding(matrices)
 
         embedding = embedding @ embedding[0].T