Updated handler.py

handler.py

@@ -7,7 +7,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 class EndpointHandler():
     def __init__(self, path=""):
         model = AutoModelForCausalLM.from_pretrained(path,
-                                                     torch_dtype=torch.float16,
+                                                     torch_dtype=torch.bfloat16,
                                                      trust_remote_code=True)
         tokenizer = AutoTokenizer.from_pretrained(path)
         #device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -19,10 +19,10 @@ class EndpointHandler():
     def __call__(self, data: Dict[str, Any]):
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
-        # with torch.autocast(self.pipeline.device.type, dtype=torch.float16):
-        outputs = self.pipeline(inputs,
-                                **parameters)
-        return outputs
+        with torch.autocast(self.pipeline.device.type, dtype=torch.bfloat16):
+            outputs = self.pipeline(inputs,
+                                    **parameters)
+            return outputs
 
 
 # class EndpointHandler: