diff --git a/sc2ts/utils.py b/sc2ts/utils.py
index 6b16f85..353da7a 100644
--- a/sc2ts/utils.py
+++ b/sc2ts/utils.py
@@ -365,7 +365,7 @@ def sort_mutation_label(s):
                 return np.inf  # put at the end
 
     if ti is None:
-        ti = TreeInfo(ts)
+        ti = sc2ts.info.TreeInfo(ts)
     if node_size is None:
         node_size = 2800
     if edge_font_size is None: