diff --git a/src/moscot/problems/time/_mixins.py b/src/moscot/problems/time/_mixins.py
index 4673934c1..0fde8a334 100644
--- a/src/moscot/problems/time/_mixins.py
+++ b/src/moscot/problems/time/_mixins.py
@@ -379,7 +379,7 @@ def sankey(
             cell_transitions_updated = cell_transitions
 
         if threshold is not None:
-            for ct in cell_transitions:
+            for ct in cell_transitions_updated:
                 ct[ct < threshold] = 0.0
 
         if key_added is None: