Update handler.py

handler.py

@@ -5,7 +5,7 @@ from io import BytesIO
 from transformers import AutoModel, AutoTokenizer
 
 class EndpointHandler:
-    def __init__(self, path=""):
+    def __init__(self, path="/repository"):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
         # Load the model
@@ -14,7 +14,8 @@ class EndpointHandler:
             trust_remote_code=True,
             attn_implementation='sdpa',
             torch_dtype=torch.bfloat16 if self.device.type == "cuda" else torch.float32,
-        ).to(self.device)
+            device_map="auto",
+        )
         self.model.eval()
         
         # Load the tokenizer