diff --git a/official/nlp/modeling/layers/__init__.py b/official/nlp/modeling/layers/__init__.py
index b94404c5382efc1de8f75b500836c99be22fa57f..0f0e17152722d11c03258ad37e6dd1565b7bc2d6 100644
--- a/official/nlp/modeling/layers/__init__.py
+++ b/official/nlp/modeling/layers/__init__.py
@@ -17,6 +17,8 @@
 They can be used to assemble new `tf.keras` layers or models.
 """
 # pylint: disable=wildcard-import
+
+from official.nlp.modeling.layers import util
 from official.nlp.modeling.layers.attention import *
 from official.nlp.modeling.layers.bigbird_attention import BigBirdAttention
 from official.nlp.modeling.layers.bigbird_attention import BigBirdMasks