System Info

Transformers 4.44
PyTorch version: 2.3.1
Ubuntu 24.04
A100 GPU

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Almost the exact same sample code found on https://huggingface.co/ai21labs/Jamba-v0.1 The only modification is device_map

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=216)

print(tokenizer.batch_decode(outputs)

Expected behavior

The error is thrown upon trying to perform inference. The easiest fix I have seen is by going to 'modeling_jamba.py' and making HybridMambaAttentionDynamicCache inherit from 'nn.Module' as such:

class HybridMambaAttentionDynamicCache(DynamicCache, nn.Module):

Then adding the superclass intitialization to the '__init__' method:

def __init__(self, config, batch_size, dtype=torch.float16, device=None):
    super(HybridMambaAttentionDynamicCache, self).__init__()