Pretrained models for the repo: https://github.com/giddyyupp/wait and paper https://arxiv.org/abs/2310.04901