(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None)
| 166 | |
| 167 | class ResidualAttentionBlock(nn.Module): |
| 168 | def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): |
| 169 | super().__init__() |
| 170 | |
| 171 | self.attn = MultiheadAttention(d_model, n_head) |
| 172 | self.ln_1 = LayerNorm(d_model) |
| 173 | self.mlp = nn.Sequential(OrderedDict([ |
| 174 | ("c_fc", nn.Linear(d_model, d_model * 4)), |
| 175 | ("gelu", QuickGELU()), |
| 176 | ("c_proj", nn.Linear(d_model * 4, d_model)) |
| 177 | ])) |
| 178 | self.ln_2 = LayerNorm(d_model) |
| 179 | self.attn_mask = attn_mask |
| 180 | |
| 181 | self.attn_probs = None |
| 182 | self.attn_grad = None |
| 183 | |
| 184 | def set_attn_probs(self, attn_probs): |
| 185 | self.attn_probs = attn_probs |
nothing calls this directly
no test coverage detected