first

app.py

@@ -7,11 +7,15 @@ import torch
 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')
 
 pipeline = transformers.pipeline(
     "text-generation",
-    model=model_name,
-    model_kwargs={"torch_dtype": torch.bfloat16},
+    model=model,
+    tokenizer=tokenizer,
     device="cuda",
 )