HemanM commited on
Commit
da84b54
·
verified ·
1 Parent(s): 20a3b86

Update evo_model.py

Browse files
Files changed (1) hide show
  1. 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, 2) # ✅ Matches saved model exactly
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]