diff --git a/export/orbax/export/tensorflow_export.py b/export/orbax/export/tensorflow_export.py
index 8470beec4..0feb7ab3d 100644
--- a/export/orbax/export/tensorflow_export.py
+++ b/export/orbax/export/tensorflow_export.py
@@ -26,6 +26,8 @@
 from orbax.export import utils
 import tensorflow as tf
 
+from .security.bcid.tree_verity.public import tree_verity_pb2
+
 obx_export_config = config.config
 maybe_reraise = reraise_utils.maybe_reraise