map rag model to device in eval rag

arcee-ai · Nov 23, 2023 · a847640 · a847640
1 parent e6c3d29
commit a847640
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/dalm/eval/eval_rag.py b/dalm/eval/eval_rag.py
@@ -196,6 +196,11 @@ def evaluate_rag(
     )
     # peft config and wrapping
     rag_model.attach_pre_trained_peft_layers(retriever_peft_model_path, generator_peft_model_path, device)
+    # mapping rag retriever and generator model to device
+    if retriever_peft_model_path is None:
+      rag_model.retriever_model.eval().to(device)
+    if generator_peft_model_path is None:
+        rag_model.generator_model.eval().to(device)
     unique_passage_dataset, passage_embeddings_array = get_passage_embeddings(
         processed_datasets,
         passage_column_name,