diff --git a/vmoe/data/input_pipeline.py b/vmoe/data/input_pipeline.py
index 4ffda0f..cec5ffa 100644
--- a/vmoe/data/input_pipeline.py
+++ b/vmoe/data/input_pipeline.py
@@ -20,6 +20,7 @@
 import ast
 from typing import Any, Callable, Dict, Optional, Union
 
+from absl import logging
 import clu.data
 import jax
 import ml_collections