Update app.py

app.py

@@ -16,23 +16,79 @@ st.set_page_config(
 
 MODEL_ID = "dejanseo/query-grounding"
 HF_TOKEN = os.getenv("HF_TOKEN")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# --- FIX: avoid meta tensors, force CPU load first with full weights ---
-model = AutoModelForSequenceClassification.from_pretrained(
-    MODEL_ID,
-    token=HF_TOKEN,
-    low_cpu_mem_usage=False,      # ensure full materialization
-    torch_dtype=torch.float32     # avoid meta tensors
-)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
-# safe device move
-model.to(device)
-model.eval()
+PREFERRED_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _has_meta_params(m: torch.nn.Module) -> bool:
+    for p in m.parameters():
+        if getattr(p, "is_meta", False):
+            return True
+    return False
+
+
+def _first_real_param_device(m: torch.nn.Module) -> torch.device:
+    for p in m.parameters():
+        if not getattr(p, "is_meta", False):
+            return p.device
+    return torch.device("cpu")
+
+
+@st.cache_resource(show_spinner=False)
+def load_model_and_tokenizer():
+    # Attempt 1: normal full load (no meta), then move to preferred device
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_ID,
+        token=HF_TOKEN,
+        low_cpu_mem_usage=False,
+        torch_dtype="auto",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+
+    # If anything is still meta, fallback to device_map loading (do NOT call .to() after that)
+    if _has_meta_params(model):
+        if torch.cuda.is_available():
+            model = AutoModelForSequenceClassification.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                torch_dtype="auto",
+                device_map="auto",
+            )
+        else:
+            # CPU fallback retry without dtype hint
+            model = AutoModelForSequenceClassification.from_pretrained(
+                MODEL_ID,
+                token=HF_TOKEN,
+                low_cpu_mem_usage=False,
+            )
+
+    # Only call .to() if the model is not dispatched by Accelerate/device_map
+    if not hasattr(model, "hf_device_map"):
+        if _has_meta_params(model):
+            raise RuntimeError(
+                "Model parameters are still on the meta device after loading. "
+                "This is usually a torch/transformers/accelerate version or memory/offload issue."
+            )
+        model.to(PREFERRED_DEVICE)
+
+    model.eval()
+    return model, tokenizer
+
+
+model, tokenizer = load_model_and_tokenizer()
+
 
 def classify(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
+    exec_device = _first_real_param_device(model)
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    )
+    inputs = {k: v.to(exec_device) for k, v in inputs.items()}
+
     with torch.no_grad():
         logits = model(**inputs).logits
         probs = torch.softmax(logits, dim=-1).squeeze().cpu()
@@ -40,6 +96,7 @@ def classify(prompt: str):
         confidence = probs[pred].item()
         return pred, confidence
 
+
 # Font and style overrides
 st.markdown("""
     <style>