Spaces:
Sleeping
Sleeping
Update evo_model.py
Browse files- evo_model.py +2 -2
evo_model.py
CHANGED
@@ -36,9 +36,9 @@ class EvoTransformerV22(nn.Module):
|
|
36 |
super().__init__()
|
37 |
self.encoder = EvoEncoder(d_model=512, num_heads=8, ffn_dim=1024, num_layers=6, memory_enabled=True)
|
38 |
self.pool = nn.AdaptiveAvgPool1d(1)
|
39 |
-
self.classifier = nn.Linear(512,
|
40 |
|
41 |
def forward(self, input_ids):
|
42 |
x = self.encoder(input_ids)
|
43 |
x = self.pool(x.transpose(1, 2)).squeeze(-1)
|
44 |
-
return self.classifier(x)
|
|
|
36 |
super().__init__()
|
37 |
self.encoder = EvoEncoder(d_model=512, num_heads=8, ffn_dim=1024, num_layers=6, memory_enabled=True)
|
38 |
self.pool = nn.AdaptiveAvgPool1d(1)
|
39 |
+
self.classifier = nn.Linear(512, 1) # ✅ Matches checkpoint
|
40 |
|
41 |
def forward(self, input_ids):
|
42 |
x = self.encoder(input_ids)
|
43 |
x = self.pool(x.transpose(1, 2)).squeeze(-1)
|
44 |
+
return self.classifier(x) # Output: [batch_size, 1]
|