diff --git a/serving/tools/timeline/gen_timeline.py b/serving/tools/timeline/gen_timeline.py
index f055e473fa0..d56c1b39897 100644
--- a/serving/tools/timeline/gen_timeline.py
+++ b/serving/tools/timeline/gen_timeline.py
@@ -1,6 +1,6 @@
 import sys
-import config_pb2
-import timeline
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import timeline
 
 def gen_timeline(src_name, dest_name):
   run_metadata = config_pb2.RunMetadata()