Update handler.py

handler.py

@@ -9,8 +9,16 @@ class EndpointHandler:
         # load model and tokenizer from path
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
-            path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True
+            path,
+            return_dict=True,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True
         )
+        
+        # self.model = AutoModelForCausalLM.from_pretrained(
+        #     path, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True
+        # )
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]: