| 228 | |
| 229 | # Copied from transformers.models.gptj.modeling_gptj.GPTJMLP with GPTJ->Moss |
| 230 | class MossMLP(nn.Module): |
| 231 | def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * embed_dim |
| 232 | super().__init__() |
| 233 | embed_dim = config.n_embd |
| 234 | |
| 235 | self.fc_in = nn.Linear(embed_dim, intermediate_size) |
| 236 | self.fc_out = nn.Linear(intermediate_size, embed_dim) |
| 237 | |
| 238 | self.act = ACT2FN[config.activation_function] |
| 239 | self.dropout = nn.Dropout(config.resid_pdrop) |
| 240 | |
| 241 | def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTensor: |
| 242 | hidden_states = self.fc_in(hidden_states) |
| 243 | hidden_states = self.act(hidden_states) |
| 244 | hidden_states = self.fc_out(hidden_states) |
| 245 | hidden_states = self.dropout(hidden_states) |
| 246 | return hidden_states |
| 247 | |
| 248 | |
| 249 | # Copied from transformers.models.gptj.modeling_gptj.GPTJBlock with GPTJ->Moss |