Fix device mapping for CPU/GPU compatibility

- Remove device_map='auto' to avoid meta tensor issues
- Explicitly handle CPU vs GPU device selection
- Use bfloat16 for CPU, float16 for GPU
- Fix device reference in chat function

app.py

@@ -12,15 +12,21 @@ ADAPTER_MODEL = "robworks-software/ccisd-teks-educator-mistral7b"
 
 print(f"📥 Loading base model: {BASE_MODEL}...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+# Check if GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"   Using device: {device}")
+
+# Load base model (use bfloat16 for CPU compatibility)
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    dtype=torch.float16,
-    device_map="auto",
+    torch_dtype=torch.bfloat16 if device == "cpu" else torch.float16,
     low_cpu_mem_usage=True
 )
 
 print(f"🔧 Loading LoRA adapter: {ADAPTER_MODEL}...")
 model = PeftModel.from_pretrained(model, ADAPTER_MODEL)
+model = model.to(device)
 
 print("✅ Model loaded successfully!")
 
@@ -47,7 +53,7 @@ def chat(message, history):
     prompt += f"[INST] {system_message}\n\n{message} [/INST]"
 
     # Tokenize
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     # Generate response
     with torch.no_grad():