diff --git a/example/test_helm_d.sh b/example/test_helm_d.sh
index 81230e1..e54a13d 100644
--- a/example/test_helm_d.sh
+++ b/example/test_helm_d.sh
@@ -1,9 +1,11 @@
 lm-eval \
   --model helm_d_115M \
+  --model_args ckpt_dir=... \
   --tasks commonsense_qa,openbookqa,hellaswag \
   --num_fewshot 0
 
 lm-eval \
   --model helm_d_115M \
+  --model_args ckpt_dir=... \
   --tasks mmlu,arc_challenge \
   --num_fewshot 5
\ No newline at end of file
diff --git a/example/test_mice_120M.sh b/example/test_mice_120M.sh
index fac2c67..9441166 100644
--- a/example/test_mice_120M.sh
+++ b/example/test_mice_120M.sh
@@ -1,9 +1,11 @@
 lm-eval \
   --model helm_mice_120M \
+  --model_args ckpt_dir=... \
   --tasks commonsense_qa,openbookqa,hellaswag \
   --num_fewshot 0
 
 lm-eval \
   --model helm_mice_120M \
+  --model_args ckpt_dir=... \
   --tasks mmlu,arc_challenge \
   --num_fewshot 5
\ No newline at end of file
diff --git a/example/test_mice_1B.sh b/example/test_mice_1B.sh
index a203472..2182627 100644
--- a/example/test_mice_1B.sh
+++ b/example/test_mice_1B.sh
@@ -1,9 +1,11 @@
 lm-eval \
   --model helm_mice_1b \
+  --model_args ckpt_dir=... \
   --tasks commonsense_qa,openbookqa,hellaswag \
   --num_fewshot 0
 
 lm-eval \
   --model helm_mice_1b \
+  --model_args ckpt_dir=... \
   --tasks mmlu,arc_challenge \
   --num_fewshot 5
\ No newline at end of file
diff --git a/lm-evaluation-harness/lm_eval/models/hypercore/models/graph_retriever.py b/lm-evaluation-harness/lm_eval/models/hypercore/models/graph_retriever.py
index cb8f666..4c6f289 100644
--- a/lm-evaluation-harness/lm_eval/models/hypercore/models/graph_retriever.py
+++ b/lm-evaluation-harness/lm_eval/models/hypercore/models/graph_retriever.py
@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from torch_geometric.nn.nlp.llm import BOS, LLM, MAX_NEW_TOKENS
+from torch_geometric.llm.models.llm import BOS, LLM, MAX_NEW_TOKENS
 from torch_geometric.utils import scatter
 from ..nn.peft import *
 
diff --git a/setup_training.sh b/setup_training.sh
index e2f8e03..057b2a4 100644
--- a/setup_training.sh
+++ b/setup_training.sh
@@ -3,4 +3,7 @@ pip install torch_geometric
 pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.6.0+cu124.html
 pip install -r requirements.txt
 pip install scipy==1.13.1 numpy==1.26.4
-pip install llm-foundry==0.18.0
\ No newline at end of file
+pip install llm-foundry==0.18.0
+cd lm-evaluation-harness
+pip install -e .
+cd ..