Updated handler.py

handler.py

@@ -7,11 +7,11 @@ class EndpointHandler:
     def __init__(self, path=""):
         # load model and tokenizer from path
         self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model = AutoModelForCausalLM.from_pretrained(path,
-                                                          device_map="auto",
+                                                          # device_map="auto",
                                                           torch_dtype=torch.float16,
-                                                          trust_remote_code=True)
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+                                                          trust_remote_code=True).to(self.device)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         # process input