Updated handler.py

handler.py

@@ -8,13 +8,14 @@ class EndpointHandler():
     def __init__(self, path=""):
         model = AutoModelForCausalLM.from_pretrained(path,
                                                      torch_dtype=torch.bfloat16,
-                                                     trust_remote_code=True)
+                                                     trust_remote_code=True,
+                                                     device_map="auto")
+        print(model.hf_device_map)
         tokenizer = AutoTokenizer.from_pretrained(path)
         #device = "cuda:0" if torch.cuda.is_available() else "cpu"
         self.pipeline = transformers.pipeline('text-generation',
                                               model=model,
-                                              tokenizer=tokenizer,
-                                              device_map="auto")
+                                              tokenizer=tokenizer)
 
     def __call__(self, data: Dict[str, Any]):
         inputs = data.pop("inputs", data)