Update handler.py

handler.py

@@ -11,7 +11,6 @@ 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]:
@@ -19,18 +18,16 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
 
-        with torch.autocast('cuda'):
-            # preprocess
-            inputs = self.tokenizer(inputs, return_tensors="pt").to('cuda')
-            self.model.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}]
+        # 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}]