diff --git a/src/moscot/problems/time/_mixins.py b/src/moscot/problems/time/_mixins.py
index 29a44aed0..7d030d6ec 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: