-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]
(head_mask, num_hidden_layers, dtype)
| 63 | module.weight.data.fill_(1.0) |
| 64 | |
| 65 | def _convert_head_mask_to_5d(head_mask, num_hidden_layers, dtype): |
| 66 | """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]""" |
| 67 | if head_mask.dim() == 1: |
| 68 | head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) |
| 69 | head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) |
| 70 | elif head_mask.dim() == 2: |
| 71 | head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer |
| 72 | assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}" |
| 73 | head_mask = head_mask.to(dtype=dtype) # switch to float if need + fp16 compatibility |
| 74 | return head_mask |
| 75 | |
| 76 | def get_head_mask( |
| 77 | head_mask, num_hidden_layers: int, |