Cool repo, fyi there is a new torch api https://github.com/pytorch/pytorch/blob/main/torch/nn/attention/flex_attention.py
Cool repo, fyi there is a new torch api https://github.com/pytorch/pytorch/blob/main/torch/nn/attention/flex_attention.py