Update model.py

model.py

@@ -1,11 +1,14 @@
+import torch
 import torch.nn as nn
 
 class SimpleEvoModel(nn.Module):
-    def __init__(self, input_dim=768, hidden_dim=256, output_dim=2):
+    def __init__(self, input_dim=384, hidden_dim=256, output_dim=2, dropout=0.1):
         super().__init__()
         self.model = nn.Sequential(
+            nn.LayerNorm(input_dim),
             nn.Linear(input_dim, hidden_dim),
             nn.ReLU(),
+            nn.Dropout(dropout),
             nn.Linear(hidden_dim, output_dim)
         )