upd

app.py

@@ -4,8 +4,12 @@ from transformers import LlamaTokenizer
 from transformers import LlamaForCausalLM, GenerationConfig
 from peft import PeftModel
 import torch
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
 
-
+device_map={'': 0}
 def generate_instruction_prompt(instruction, input=None):
     if input:
         return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@@ -45,7 +49,7 @@ def evaluate(
     )
     prompt = generate_instruction_prompt(instruction, input)
     inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"].cuda()
+    input_ids = inputs["input_ids"].to(device)
     generation_output = model.generate(
         input_ids=input_ids,
         generation_config=generation_config,
@@ -64,19 +68,19 @@ def load_lora(lora_path, base_model="decapoda-research/llama-7b-hf"):
     model = LlamaForCausalLM.from_pretrained(
         base_model,
         # load_in_8bit=True,
-        device_map="auto",
+        device_map=device_map,
         low_cpu_mem_usage=True,
     )
     lora = PeftModel.from_pretrained(
         model,
         lora_path,
-        device_map="auto",
+        device_map=device_map,
     )
     return lora
 
 
 base_model = "decapoda-research/llama-13b-hf"
-tokenizer = LlamaTokenizer.from_pretrained(base_model)
+tokenizer = LlamaTokenizer.from_pretrained(base_model, device_map=device_map)
 # question = "如果今天是星期五, 那么后天是星期几?"
 model = load_lora(lora_path="facat/alpaca-lora-cn-13b", base_model=base_model)