Update app.py

app.py

@@ -8,7 +8,7 @@ from theme_dropdown import create_theme_dropdown
 
 model_name = "RootYuan/RootYuan-RedLing-7B-v0.1"
 max_new_tokens = 2048
-device = 'cpu'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 DEFAULT_SYSTEM_MESSAGE = """
@@ -22,7 +22,7 @@ EOT_TOKEN = "<EOT>"
 PROMPT_TEMPLATE = "USER:{user}<EOT>ASSISTANT:{assistant}{eos_token}"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
 
 dropdown, js = create_theme_dropdown()