(self, heads)
| 285 | self.pruned_heads = set() |
| 286 | |
| 287 | def prune_heads(self, heads): |
| 288 | if len(heads) == 0: |
| 289 | return |
| 290 | heads, index = find_pruneable_heads_and_indices( |
| 291 | heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads |
| 292 | ) |
| 293 | |
| 294 | # Prune linear layers |
| 295 | self.self.query = prune_linear_layer(self.self.query, index) |
| 296 | self.self.key = prune_linear_layer(self.self.key, index) |
| 297 | self.self.value = prune_linear_layer(self.self.value, index) |
| 298 | self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) |
| 299 | |
| 300 | # Update hyper params and store pruned heads |
| 301 | self.self.num_attention_heads = self.self.num_attention_heads - len(heads) |
| 302 | self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads |
| 303 | self.pruned_heads = self.pruned_heads.union(heads) |
| 304 | |
| 305 | def forward( |
| 306 | self, |
no test coverage detected