Update app.py

app.py

@@ -46,7 +46,7 @@ PROMPT_DICT = {
 }
 from llama_rope_scaled_monkey_patch import replace_llama_rope_with_scaled_rope
 replace_llama_rope_with_scaled_rope()
-model = transformers.AutoModelForCausalLM.from_pretrained(
+base_model = transformers.AutoModelForCausalLM.from_pretrained(
             base_model,
             torch_dtype=torch.float16,
             cache_dir=cache_dir,
@@ -54,12 +54,15 @@ model = transformers.AutoModelForCausalLM.from_pretrained(
         )
 
 model = PeftModel.from_pretrained(
-            model,
+            base_model,
             lora_weights,
             device_map="auto",
             cache_dir=cache_dir,
             torch_dtype=torch.float16,
+            assign=True
         )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 tokenizer =  AutoTokenizer.from_pretrained(base_model,use_fast=False,cache_dir=cache_dir)
 tokenizer.pad_token = tokenizer.unk_token
 def generate_prompt(instruction, input=None):