diff --git a/examples/python/jobshop_ft06_sat.py b/examples/python/jobshop_ft06_sat.py
index 340e2a7507..c9f481443b 100644
--- a/examples/python/jobshop_ft06_sat.py
+++ b/examples/python/jobshop_ft06_sat.py
@@ -13,6 +13,7 @@
 
 import collections
 from ortools.sat.python import cp_model
+from ortools.sat.python import visualization
 
 
 def main():
diff --git a/ortools/sat/python/visualization.py b/ortools/sat/python/visualization.py
index 25fd9a1079..8fff09863e 100644
--- a/ortools/sat/python/visualization.py
+++ b/ortools/sat/python/visualization.py
@@ -13,19 +13,15 @@
 """Collection of helpers to visualize cp_model solutions in colab."""
 
 import random
-from IPython.display import display
-from IPython.display import SVG
-import plotly.figure_factory as ff
-import plotly.offline as pyo
-import svgwrite
-
-
-def RunFromIPython():
-    try:
-        __IPYTHON__
-        return True
-    except NameError:
-        return False
+try:
+  from IPython.display import display
+  from IPython.display import SVG
+  import plotly.figure_factory as ff
+  import plotly.offline as pyo
+  import svgwrite
+  run_from_ipython = True
+except ImportError:
+  run_from_ipython = False
 
 
 def ToDate(v):
@@ -164,4 +160,3 @@ class SvgWrapper(object):
         font_family='sans-serif',
         font_size='%dpx' % (self.__scaling / 2))
     self.__dwg.add(text)
-