potential fix for an issue in generate, identified by @Jiang-Stan https://github.com/lucidrains/soundstorm-pytorch/issues/28