Update app.py

app.py

@@ -9,8 +9,8 @@ import time
 model_id = "HuggingFaceTB/SmolLM-1.7B"
 assistant_id = "HuggingFaceTB/SmolLM-135M"
 
-model = AutoModelForCausalLM.from_pretrained(model_id)  #.to(dtype=torch.bfloat16, device="cuda")
-assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id)  #.to(dtype=torch.bfloat16, device="cuda")
+model = AutoModelForCausalLM.from_pretrained(model_id).to(dtype=torch.bfloat16, device="cuda")
+assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(dtype=torch.bfloat16, device="cuda")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 @spaces.GPU
@@ -56,9 +56,9 @@ def reset_textbox():
 with gr.Blocks() as demo:
     gr.Markdown(
         "# 🤗 Assisted Generation Demo\n"
-        f"- Model: {model_id}\n"
-        f"- Assistant Model: {assistant_id}\n"
-        "- Running on CPU"
+        f"- Model: {model_id} (BF16)\n"
+        f"- Assistant Model: {assistant_id} (BF16)\n"
+        "- Running on a T4 GPU"
     )
 
     with gr.Row():