Update handler.py

handler.py

@@ -11,6 +11,7 @@ class EndpointHandler:
         self.model = AutoModelForCausalLM.from_pretrained(
             path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True
         )
+        print('===> cuda.is_available', torch.cuda.is_available())
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
@@ -18,16 +19,17 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
 
-        # preprocess
-        inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device)
-
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            outputs = self.model.generate(**inputs, **parameters)
-        else:
-            outputs = self.model.generate(**inputs)
-
-        # postprocess the prediction
-        prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        return [{"generated_text": prediction}]
+        with torch.autocast('cuda'):
+            # preprocess
+            inputs = self.tokenizer(inputs, return_tensors="pt").to('cuda')
+    
+            # pass inputs with all kwargs in data
+            if parameters is not None:
+                outputs = self.model.generate(**inputs, **parameters)
+            else:
+                outputs = self.model.generate(**inputs)
+    
+            # postprocess the prediction
+            prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+            return [{"generated_text": prediction}]