Setting MPS backend causes mamba-minimal to generate gibberish and crash on M1 Mac

[oselivanov]

Trying to run it on MPS backend on Mac M1 Max 64Gb modifying demo.ipynb like this:

---

Changes to notebook:

device = 'mps'

model = Mamba.from_pretrained(pretrained_model_name).to(device)

input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)

---

Getting gibberish and dead python kernel in dozen of tokens:

for t in generate(model, tokenizer, 'Mamba is the'): print(t, end='')

ocardial goal fibrobl ( triglycer ા judgement extraordinacebook CURIAM

and then it crashes.