Add explicit error if empty batch received (#2115)

* Add explicit error if empty batch received * Move error check to cover all empty iterables
huggingface · Nov 3, 2023 · 4f10031 · 4f10031
1 parent fa6f430
commit 4f10031
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
@@ -231,6 +231,9 @@ def find_batch_size(data):
     Returns:
         `int`: The batch size.
     """
+    if isinstance(data, (tuple, list, Mapping)) and (len(data) == 0):
+        raise ValueError(f"Cannot find the batch size from empty {type(data)}.")
+
     if isinstance(data, (tuple, list)):
         return find_batch_size(data[0])
     elif isinstance(data, Mapping):