fix mypy issue

google · Jul 19, 2023 · 68b334e · 68b334e
1 parent 50f492b
commit 68b334e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/flax/linen/recurrent.py b/flax/linen/recurrent.py
@@ -818,6 +818,7 @@ def scan_fn(
     # return_carry is True we slice the carry history and select the last valid
     # carry for each sequence. Otherwise we just use the last carry.
     if slice_carry:
+      assert seq_lengths is not None
       _, (carries, outputs) = scan_output
       # seq_lengths[None] expands the shape of the mask to match the
       # number of dimensions of the carry.