google

/

gemma-7b-it

like 1.15k

Follow

Google 6.27k

I'm running the sample code

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it", device_map="auto")

input_text = "Write me a poem about Machine Learning."
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids)
print(tokenizer.decode(outputs[0]))

but I'm getting a shape error:

RuntimeError: shape '[1, 9, 3072]' is invalid for input of size 36864