(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
)
| 140 | self.final_layer_norm = nn.LayerNorm(embed_dim) |
| 141 | |
| 142 | def forward( |
| 143 | self, |
| 144 | hidden_states: torch.Tensor, |
| 145 | attention_mask: torch.Tensor, |
| 146 | ): |
| 147 | # Attention |
| 148 | residual = hidden_states |
| 149 | hidden_states = self.self_attn_layer_norm(hidden_states) |
| 150 | hidden_states = self.self_attn(hidden_states, attention_mask) |
| 151 | hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) |
| 152 | hidden_states = residual + hidden_states |
| 153 | |
| 154 | # FFN |
| 155 | residual = hidden_states |
| 156 | hidden_states = self.final_layer_norm(hidden_states) |
| 157 | hidden_states = F.gelu(self.fc1(hidden_states)) |
| 158 | hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) |
| 159 | hidden_states = self.fc2(hidden_states) |
| 160 | hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) |
| 161 | hidden_states = residual + hidden_states |
| 162 | return hidden_states |
| 163 | |
| 164 | def forward_chunk( |
| 165 | self, |
nothing calls this directly
no outgoing calls
no test coverage detected