(
head_mask, num_hidden_layers: int,
is_attention_chunked: bool = False
)
| 74 | return head_mask |
| 75 | |
| 76 | def get_head_mask( |
| 77 | head_mask, num_hidden_layers: int, |
| 78 | is_attention_chunked: bool = False |
| 79 | ): |
| 80 | if head_mask is not None: |
| 81 | head_mask = _convert_head_mask_to_5d(head_mask, num_hidden_layers, 'float16') |
| 82 | if is_attention_chunked is True: |
| 83 | head_mask = head_mask.unsqueeze(-1) |
| 84 | else: |
| 85 | head_mask = [None] * num_hidden_layers |
| 86 | |
| 87 | return head_mask |
no test coverage detected