Update app.py

app.py

@@ -1,22 +1,11 @@
 import gradio as gr
 from transformers import pipeline
-import torch
-
-try:
-    router = pipeline(
-        "text-classification",
-        model="AmirMohseni/reasoning-router-0.6b",
-        device_map="auto",
-        torch_dtype=torch.bfloat16
-    )
-except Exception as e:
-    print(f"Error loading model: {e}")
-    # Fallback to CPU if specific GPU setup fails
-    router = pipeline(
-        "text-classification",
-        model="AmirMohseni/reasoning-router-0.6b",
-    )
 
+router = pipeline(
+    "text-classification",
+    model="AmirMohseni/reasoning-router-0.6b",
+    device_map="auto",
+)
 
 # --- 2. Define the classification function ---
 def classify_prompt(prompt: str) -> dict: