I ran into the same issue and was able to resolve it. It turns out to be related to the Donut model’s MAX_TOKEN_LEN setting. My code runs successfully when MAX_TOKEN_LEN is set to 128 or lower, but the bug reappears as soon as it exceeds 128.