diff --git a/pearl/utils/scripts/cb_benchmark/cb_benchmark_config.py b/pearl/utils/scripts/cb_benchmark/cb_benchmark_config.py
index db394466..11a9925c 100644
--- a/pearl/utils/scripts/cb_benchmark/cb_benchmark_config.py
+++ b/pearl/utils/scripts/cb_benchmark/cb_benchmark_config.py
@@ -28,7 +28,7 @@
 from pearl.policy_learners.exploration_modules.contextual_bandits.ucb_exploration import (
     UCBExploration,
 )
-from pearl.test.utils import prefix_dir
+from test.utils import prefix_dir
 from pearl.utils.instantiations.environments.contextual_bandit_uci_environment import (
     SLCBEnvironment,
 )
diff --git a/pearl/utils/scripts/cb_benchmark/run_cb_benchmarks.py b/pearl/utils/scripts/cb_benchmark/run_cb_benchmarks.py
index 24319fe2..df36f767 100644
--- a/pearl/utils/scripts/cb_benchmark/run_cb_benchmarks.py
+++ b/pearl/utils/scripts/cb_benchmark/run_cb_benchmarks.py
@@ -29,7 +29,7 @@
 from pearl.replay_buffers.contextual_bandits.discrete_contextual_bandit_replay_buffer import (
     DiscreteContextualBanditReplayBuffer,
 )
-from pearl.test.utils import prefix_dir
+from test.utils import prefix_dir
 from pearl.utils.instantiations.environments.contextual_bandit_uci_environment import (
     SLCBEnvironment,
 )