RuntimeError: selected index k out of range - num patches

[aropemaker]

Hi there,

When I try and train Mulan, I am getting this error output:

Traceback (most recent call last):
File "txt2music_p1.py", line 230, in
embeds = mulan.get_audio_latents(wavs) # during training
File "/home/qowabungaqing/txt2music/lib/python3.8/site-packages/musiclm_pytorch/musiclm_pytorch.py", line 702, in get_audio_latents
audio_embeds, audio_layers = self.audio(wavs, return_all_layers = True)
File "/home/qowabungaqing/txt2music/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/home/qowabungaqing/txt2music/lib/python3.8/site-packages/musiclm_pytorch/musiclm_pytorch.py", line 479, in forward
patch_indices_keep = torch.randn(batch, n, device = device).topk(num_patches_keep, dim = -1).indices
RuntimeError: selected index k out of range

And after a bit of debugging, I found that this is what is happening when I am trying to train my dataset:
n: 680
num_patches_keep: 510
n: 0
num_patches_keep: 1

so it looks like it is getting stuck on the second patch.
I have checked that all of the wavs are not corrupt, but I don't really understand enough about what is going on to do more debugging.

Any help would be greatly appreciated, thank you.