ZekunZh · February 16, 2022 10:05 · Feb 16, 2022
diff --git a/iter_batches.py b/iter_batches.py
@@ -0,0 +1,18 @@
+def iter_batches(iterable, batch_size):
+  """Iterates over the given iterable in batches.
+
+  Args:
+      iterable: an iterable
+      batch_size: the desired batch size, or None to return the contents in
+          a single batch
+
+  Returns:
+      a generator that emits tuples of elements of the requested batch size
+      from the input
+  """
+  it = iter(iterable)
+  while True:
+    chunk = tuple(itertools.islice(it, batch_size))
+    if not chunk:
+      return
+    yield chunk
No results found