Update app.py

app.py

@@ -26,13 +26,15 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model_name = "01-ai/Yi-6B-200K"
 
 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = tokenizer.cuda()
 # tokenizer = YiTokenizer(vocab_file=model_name)
 model = transformers.AutoModelForCausalLM.from_pretrained(model_name,
-    device_map=device,
+    device_map="auto",
     torch_dtype=torch.bfloat16,
     load_in_4bit=True,
     trust_remote_code=True
 )
+model = model.cuda()
 
 # Load the model and tokenizer using transformers
 # model = AutoModelForCausalLM.from_pretrained("01-ai/Yi-6B-200K", trust_remote_code=True)