From a77a34759fc974b19feba5f2e0af9400a3830f20 Mon Sep 17 00:00:00 2001
From: winter <2453101190@qq.com>
Date: Thu, 4 Jul 2024 19:54:18 +0800
Subject: [PATCH] use default gpu 0

---
 demo/hf_based_demo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/demo/hf_based_demo.py b/demo/hf_based_demo.py
index 75df156..4eaa184 100644
--- a/demo/hf_based_demo.py
+++ b/demo/hf_based_demo.py
@@ -34,7 +34,7 @@ else:
 # init model and tokenizer
 path = args.model_path
 tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, device_map="auto", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch_dtype, device_map="cuda:0", trust_remote_code=True)
 
 model_architectures = model.config.architectures[0]