HemanM commited on
Commit
2dccd99
·
verified ·
1 Parent(s): 3d17dd0

Update evo_model.py

Browse files
Files changed (1) hide show
  1. evo_model.py +2 -1
evo_model.py CHANGED
@@ -32,6 +32,7 @@ class FeedForward(nn.Module):
32
  self.net = nn.Sequential(
33
  nn.Linear(d_model, dim_feedforward),
34
  nn.ReLU(),
 
35
  nn.Linear(dim_feedforward, d_model)
36
  )
37
 
@@ -52,7 +53,7 @@ class TransformerBlock(nn.Module):
52
  return x
53
 
54
  class EvoDecoder(nn.Module):
55
- def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=3, dim_feedforward=1024):
56
  super().__init__()
57
  self.token_emb = nn.Embedding(vocab_size, d_model)
58
  self.pos_emb = nn.Embedding(512, d_model)
 
32
  self.net = nn.Sequential(
33
  nn.Linear(d_model, dim_feedforward),
34
  nn.ReLU(),
35
+ nn.Dropout(), # ✅ Important: was present in the training model
36
  nn.Linear(dim_feedforward, d_model)
37
  )
38
 
 
53
  return x
54
 
55
  class EvoDecoder(nn.Module):
56
+ def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=3, dim_feedforward=512):
57
  super().__init__()
58
  self.token_emb = nn.Embedding(vocab_size, d_model)
59
  self.pos_emb = nn.Embedding(512, d_model)