Update app.py

app.py

@@ -90,7 +90,7 @@ def warmup(model):
     if args.model_type == "llama-2-chat":
         prompt += " "
     input_ids = model.tokenizer([prompt]).input_ids
-    input_ids = torch.as_tensor(input_ids).to(model.device)
+    input_ids = torch.as_tensor(input_ids).to(model.base_model.device)
     for output_ids in model.ea_generate(input_ids):
         ol=output_ids.shape[1]
 @spaces.GPU(duration=30)
@@ -144,7 +144,7 @@ def bot(history, temperature, top_p, use_EaInfer, highlight_EaInfer,session_stat
         prompt += " "
 
     input_ids = model.tokenizer([prompt]).input_ids
-    input_ids = torch.as_tensor(input_ids).to(model.device)
+    input_ids = torch.as_tensor(input_ids).to(model.base_model.device)
     input_len = input_ids.shape[1]
     naive_text = []
     cu_len = input_len