try to actually use the GPU

handler.py

@@ -1,27 +1,37 @@
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
+import torch
 
 class EndpointHandler:
     def __init__(self, path=""):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
         # load the model
         tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
         model = AutoModelForCausalLM.from_pretrained(
             "Qwen/Qwen2-1.5B-Instruct",
-            torch_dtype="auto",
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
             device_map="auto"
-        )
+        ).to(device)
+        
         # create inference pipeline
-        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if device == "cuda" else -1)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
+        parameters = data.pop("parameters", {})
+
+        # Ensure inputs are on the GPU if available
+        if isinstance(inputs, str):
+            inputs = [inputs]
+
+        # Tensor input handling
+        try:
+            inputs = torch.tensor(inputs).cuda() if torch.cuda.is_available() else torch.tensor(inputs)
+        except:
+            pass  # If inputs are not tensors (e.g., strings), continue without conversion
 
         # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
-        else:
-            prediction = self.pipeline(inputs)
-        # postprocess the prediction
+        prediction = self.pipeline(inputs, **parameters)
+        
         return prediction