Update evo_model.py

evo_model.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 class EvoEncoder(nn.Module):
-    def __init__(self, d_model=384, num_heads=6, ffn_dim=1024, num_layers=6, memory_enabled=True):
+    def __init__(self, d_model=512, num_heads=8, ffn_dim=1024, num_layers=6, memory_enabled=True):
         super().__init__()
         self.embedding = nn.Embedding(30522, d_model)
         self.memory_enabled = memory_enabled
@@ -31,14 +31,13 @@ class EvoEncoder(nn.Module):
         x = self.transformer(x)
         return x
 
-
 class EvoTransformerV22(nn.Module):
     def __init__(self):
         super().__init__()
-        self.encoder = EvoEncoder(memory_enabled=True)
+        self.encoder = EvoEncoder(d_model=512, num_heads=8, ffn_dim=1024, num_layers=6, memory_enabled=True)
         self.pool = nn.AdaptiveAvgPool1d(1)
         self.classifier = nn.Sequential(
-            nn.Linear(384, 128),
+            nn.Linear(512, 128),
             nn.ReLU(),
             nn.Linear(128, 2)
         )