Update core/model_loader.py

core/model_loader.py

@@ -3,19 +3,19 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 SUPPORTED_MODELS = {
     "DistilGPT2 (Fast CPU)": "distilgpt2",
-    "TinyLlama (Better LLM)": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-    "Phi-2 (Research Heavy)": "microsoft/phi-2"
 }
 
 def load_model(model_key):
     model_name = SUPPORTED_MODELS[model_key]
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,
-        device_map="cpu"
+        dtype=torch.float32   # ✅ Use dtype instead of torch_dtype
     )
 
+    model.to("cpu")  # ✅ Explicit CPU move
     model.eval()
+
     return model, tokenizer