first

app.py

@@ -9,14 +9,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model_name = "ModularityAI/gemma-2b-datascience-it-raft"
 tokenizer_name = "google/gemma-2b-it"
 
-model = AutoModelForCausalLM.from_pretrained(model_name,torch_dtype=torch.bfloat16,device='cuda')
-tokenizer = AutoTokenizer.from_pretrained(tokenizer_name,device='cuda')
+model = AutoModelForCausalLM.from_pretrained(model_name,torch_dtype=torch.bfloat16,device_map='cuda')
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name,device_map='cuda')
 
 pipeline = transformers.pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device="cuda",
 )
 
 def format_test_question(q):