Code modified from: https://github.com/yandex-research/tab-ddpm