diff --git a/tensorflow/lite/tools/visualize.py b/tensorflow/lite/tools/visualize.py
index de7ef820079..cd4bcfa7aaf 100644
--- a/tensorflow/lite/tools/visualize.py
+++ b/tensorflow/lite/tools/visualize.py
@@ -33,7 +33,7 @@
   from tflite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
 else:
   # This file is part of tflite_runtime package.
-  from tflite_runtime import schema_py_generated as schema_fb
+  from tflite_micro.tensorflow.lite_runtime import schema_py_generated as schema_fb
 
 # A CSS description for making the visualizer
 _CSS = """