File size: 623 Bytes
d0f5e0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | """Project-local compatibility patches for the H20 CoVT environment."""
try:
import torch
import flash_attn.bert_padding as _flash_bert_padding
def _covt_pad_input(hidden_states, indices, batch, seqlen):
trailing_shape = tuple(hidden_states.shape[1:])
output = torch.zeros(
(batch * seqlen, *trailing_shape),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
output[indices] = hidden_states
return output.reshape(batch, seqlen, *trailing_shape)
_flash_bert_padding.pad_input = _covt_pad_input
except Exception:
pass
|