Update src/demo.py

src/demo.py

@@ -22,7 +22,7 @@ type2dataset = {
 
 model_id = "meta-llama/Llama-2-7b-chat-hf"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=TOKEN)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=TOKEN).eval()
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=TOKEN, device_map="auto").eval()
 
 # type2dataset = {}
 
@@ -35,7 +35,7 @@ def generate(input_text, sys_prompt) -> str:
 '''
     input_str = sys_prompt + input_text + " [/INST]"
     
-    input_ids = tokenizer(input_str, return_tensors="pt").input_ids
+    input_ids = tokenizer(input_str, return_tensors="pt").input_ids.to('cuda')
     outputs = model.generate(input_ids, max_length=512)
 
     result = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]