diff --git a/model_custom.py b/model_custom.py
index dc73f4c034d139cbdf698ac4f3478922266d408f..d2888b40567abc2e1ff1f1384cf3b6c26f75dd88 100644
--- a/model_custom.py
+++ b/model_custom.py
@@ -106,6 +106,7 @@ class Model_Common_Transformer(nn.Module):
     def forward(self, seq, charge):
         meta_ohe = torch.nn.functional.one_hot(charge - 1, self.charge_max).float()
         seq_emb = torch.nn.functional.one_hot(seq, self.nb_aa).float()
+        seq_emb = torch.permute(seq_emb, (0, 2, 1))
         print(seq_emb.size())
         emb = self.pos_embedding(self.emb(seq_emb))
         meta_enc = self.meta_enc(meta_ohe)
@@ -130,6 +131,7 @@ class Model_Common_Transformer(nn.Module):
     def forward_int(self, seq, charge):
         meta_ohe = torch.nn.functional.one_hot(charge - 1, self.charge_max).float()
         seq_emb = torch.nn.functional.one_hot(seq, self.nb_aa).float()
+        seq_emb = torch.permute(seq_emb, (0, 2, 1))
         emb = self.pos_embedding(self.emb(seq_emb))
         meta_enc = self.meta_enc(meta_ohe)
         enc = self.encoder(emb)