Update app.py: Add device=device

app.py

@@ -186,12 +186,14 @@ def on_btn_click():
 
 @st.cache_resource
 def load_model():
-    model = (AutoModelForCausalLM.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         model_name_or_path,
         trust_remote_code=True,
         use_cache=False,  # 禁用 KV 缓存
         torch_dtype=torch.bfloat16,
-        device_map="auto")).cuda()
+        device_map="auto",
+        low_cpu_mem_usage=True
+    )
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                               trust_remote_code=True)
     return model, tokenizer
@@ -279,7 +281,7 @@ def main():
                     tokenizer=tokenizer,
                     prompt=real_prompt,
                     additional_eos_token_id=92542,
-                    device='cuda:0',
+                    device=device,
                     **asdict(generation_config),
             ):
                 # Display robot response in chat message container