diff --git a/mesh_tensorflow/transformer/transformer.py b/mesh_tensorflow/transformer/transformer.py
index e3ab84f0..48a6a629 100644
--- a/mesh_tensorflow/transformer/transformer.py
+++ b/mesh_tensorflow/transformer/transformer.py
@@ -45,6 +45,7 @@
 
 import json
 import math
+import re
 
 import gin
 import mesh_tensorflow as mtf