facebook/musicgen-small · RuntimeError: probability tensor contains either `inf`, `nan` or element

When running the colab, I get the following error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-a32693f3d849> in <cell line: 3>()
      1 unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
      2 
----> 3 audio_values = model.generate(**unconditional_inputs, do_sample=True, max_new_tokens=256)

2 frames
/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py in _sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, logits_warper, **model_kwargs)
   2435             if do_sample:
   2436                 probs = nn.functional.softmax(next_token_scores, dim=-1)
-> 2437                 next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
   2438             else:
   2439                 next_tokens = torch.argmax(next_token_scores, dim=-1)

RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

I also get the same error in other applications. The error does not happen when using CPU.