Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -95,20 +95,19 @@ def test_ai_integration_http():
                 transformers_logger.setLevel(logging.ERROR)
                 
                 try:
-                    # Create pipeline
-                    # Create pipeline with vGPU device
+                    # Create pipeline with model directly on vGPU
                     pipe = pipeline(
                         "text-generation",
                         model=model_id,
-                        torch_dtype=torch.float32,  # Use full precision,
+                        model_kwargs={
+                            "torch_dtype": torch.float32,  # Use full precision
+                            "device_map": {"": device},  # Map all modules to our vGPU device
+                        },
                         use_safetensors=True,
                         trust_remote_code=True,
-                        device=device
+                        device=device  # Use our vGPU device
                     )
-                    pipe.model.eval()
-                    
-                    # Move pipeline model to vGPU
-
+                    status["pipeline_loaded"] = True
                     status['model_on_vgpu'] = True
                     
                     # Log model details