Update handler.py

handler.py

@@ -33,8 +33,8 @@ class EndpointHandler():
     def __init__(self, path=""):
         DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
         model_name = "moka-ai/m3e-base"
-        self.config = PeftConfig.from_pretrained("YingL19/gemma_10epoch_1e5_lincoln1",device_map=DEVICE)
-        self.base_model = AutoModelForCausalLM.from_pretrained("google/gemma-7b",device_map=DEVICE)
+        self.config = PeftConfig.from_pretrained("YingL19/gemma_10epoch_1e5_lincoln1")
+        self.base_model = AutoModelForCausalLM.from_pretrained("google/gemma-7b",device_map="auto")
         self.newmodel = PeftModel.from_pretrained(self.base_model, "YingL19/gemma_10epoch_1e5_lincoln1",device_map=DEVICE)
         self.tokenizer = AutoTokenizer.from_pretrained("YingL19/gemma_10epoch_1e5_lincoln1",device_map=DEVICE)