diff --git a/efficientdet/dataloader_test.py b/efficientdet/dataloader_test.py
index 066ff9d0f..a4d2e4ff0 100644
--- a/efficientdet/dataloader_test.py
+++ b/efficientdet/dataloader_test.py
@@ -18,7 +18,7 @@
 
 import dataloader
 import hparams_config
-from brain_automl.efficientdet import test_util
+import test_util
 
 from keras import anchors
 from object_detection import tf_example_decoder
diff --git a/efficientdet/keras/inference_test.py b/efficientdet/keras/inference_test.py
index 47edce955..e682d447b 100644
--- a/efficientdet/keras/inference_test.py
+++ b/efficientdet/keras/inference_test.py
@@ -17,7 +17,7 @@
 import tempfile
 from absl import logging
 import tensorflow as tf
-from brain_automl.efficientdet import test_util
+import test_util
 from keras import efficientdet_keras
 from keras import inference