Update app.py

app.py

@@ -10,27 +10,38 @@ def chat_with_model(message, history, max_tokens, temperature):
     try:
         print("Loading model...")
         
-        # Load model and tokenizer - let ZeroGPU handle device placement
+        # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+        
+        # Load model WITHOUT device_map to avoid CPU placement
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
-            device_map="auto"  # Let transformers handle GPU placement
+            # Remove device_map="auto" - it's causing CPU placement
         )
         
+        # AFTER loading, move to GPU explicitly
+        if torch.cuda.is_available():
+            model = model.to('cuda')
+            print(f"✅ Model moved to GPU: {next(model.parameters()).device}")
+        else:
+            print("❌ No GPU available")
+        
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
-        print(f"Model loaded successfully on device: {next(model.parameters()).device}")
-        
         # Simple prompt format
         prompt = f"User: {message}\nAssistant:"
         
-        # Tokenize - let the model handle device placement
+        # Tokenize and move inputs to same device as model
         inputs = tokenizer(prompt, return_tensors="pt")
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         
-        # Generate - the model will automatically handle device placement
+        print(f"✅ Inputs on device: {inputs['input_ids'].device}")
+        
+        # Generate
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
@@ -44,6 +55,8 @@ def chat_with_model(message, history, max_tokens, temperature):
         # Decode response
         response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
         
+        print(f"✅ Generated: {response[:50]}...")
+        
         # Update history
         history.append([message, response])
         return history, history, ""