diff --git a/grain/_src/python/dataset/transformations/interleave.py b/grain/_src/python/dataset/transformations/interleave.py
index 797460dd..e0f01e38 100644
--- a/grain/_src/python/dataset/transformations/interleave.py
+++ b/grain/_src/python/dataset/transformations/interleave.py
@@ -214,7 +214,7 @@ def make_source(filename: str) -> grain.IterDataset:
 
   def __init__(
       self,
-      datasets: Sequence[dataset.IterDataset[T] | dataset.MapDataset[T]],
+      datasets: dataset.MapDataset[T] | Sequence[dataset.IterDataset[T] | dataset.MapDataset[T]],
       *,
       cycle_length: int,
       num_make_iter_threads: int = 1,