Update README.md

README.md

@@ -34,8 +34,8 @@ tokenizer = AutoTokenizer.from_pretrained("GeneZC/MiniChat-3B", use_fast=False)
 # GPU.
 model = AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="auto", torch_dtype=torch.float16).eval()
 # CPU.
-# model = AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="cpu", torch_dtype=torch.float16).eval()
-
+# model = AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="cpu", torch_dtype=torch.float32).eval()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 conv = get_default_conv_template("minichat")
 
 question = "Implement a program to find the common elements in two arrays without using any extra data structures."
@@ -44,7 +44,7 @@ conv.append_message(conv.roles[1], None)
 prompt = conv.get_prompt()
 input_ids = tokenizer([prompt]).input_ids
 output_ids = model.generate(
-    torch.as_tensor(input_ids).cuda(),
+    torch.as_tensor(input_ids).to(device),
     do_sample=True,
     temperature=0.7,
     max_new_tokens=1024,