diff --git a/ffn/jax/input_pipeline.py b/ffn/jax/input_pipeline.py
index 2dc8fc6..7d08e78 100644
--- a/ffn/jax/input_pipeline.py
+++ b/ffn/jax/input_pipeline.py
@@ -28,6 +28,7 @@
 from ffn.training import examples
 from ffn.training import inputs
 from ffn.training import model as ffn_model
+from google.protobuf import text_format
 import jax
 import jmp
 import ml_collections
diff --git a/ffn/training/inputs.py b/ffn/training/inputs.py
index 1efed9e..3eba712 100644
--- a/ffn/training/inputs.py
+++ b/ffn/training/inputs.py
@@ -26,6 +26,7 @@
 from connectomics.volume import metadata
 from ffn.training import augmentation
 from ffn.training import variables
+from google.protobuf import text_format
 import numpy as np
 import tensorflow.compat.v1 as tf
 from tensorflow.io import gfile
diff --git a/ffn/training/inputs_test.py b/ffn/training/inputs_test.py
index 14a0778..e51108c 100644
--- a/ffn/training/inputs_test.py
+++ b/ffn/training/inputs_test.py
@@ -22,6 +22,7 @@
 from connectomics.common import tuples
 from connectomics.volume import metadata
 from ffn.training import inputs
+from google.protobuf import text_format
 import numpy as np
 import tensorflow.compat.v1 as tf