Update handler.py

handler.py

@@ -34,7 +34,7 @@ class EndpointHandler():
         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")
+        self.base_model = AutoModelForCausalLM.from_pretrained("google/gemma-7b",device_map=DEVICE)
         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)
 
@@ -68,7 +68,7 @@ ASSISTANT:'''
             template=self.prompt_template,
         )
 
-        embeddings = HuggingFaceEmbeddings(model_name=model_name)
+        embeddings = HuggingFaceEmbeddings(model_name=model_name,model_kwargs={'device':DEVICE})
         embeddings.client = SentenceTransformer(model_name, device=DEVICE)
         ragdb = Chroma(persist_directory="/Lincoln_DB", embedding_function=embeddings)