Repository: huggingface/accelerate

Issue:
Improve error message in find_batch_size()
Function [find_batch_size()](https://github.com/huggingface/accelerate/blob/55747318a0f47cdfbc281e11269ba96214e4092d/src/accelerate/utils/operations.py#L224) produces unclear error when it gets empty dict `{}` as a parameter:
```
      239 elif not isinstance(data, np.ndarray):  #torch.Tensor):
      240     raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
----> 241 return data.shape[0]

AttributeError: 'dict' object has no attribute 'shape'
```

For example, this happens when training in `transformer` library is not configured well, and training parameter `remove_unused_columns=True`. In this case `transformer` [can remove all data columns from the dataset](https://github.com/huggingface/transformers/blob/acc394c4f5e1283c19783581790b3dc3105a3697/src/transformers/trainer.py#L752) during the data preparation. This data then goes to `find_batch_size()` ends up in the mentioned error with dictionary.

From the source code it is hard to realize that error happened because of empty batch was passed to this function (because the whole dataset was empty). It rather signals that dataset has improper format, which makes investigation much longer.

In my opinion, the problem come from the [following line](https://github.com/huggingface/accelerate/blob/5b3f3b99d6aa1d71eaa3380af87a6ca6126505fd/src/accelerate/utils/operations.py#L237):
```
    if isinstance(data, (tuple, list)):
        return find_batch_size(data[0])
    elif isinstance(data, Mapping):
        for k in data.keys():   # <== looks like loop, but it is not, just a quick way to get key from dict; when dict is empty, we unexpectedly go to "return data.shape[0]"
            return find_batch_size(data[k])
    elif not isinstance(data, torch.Tensor):
        raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
    return data.shape[0]
```

I see the following ways how we can make it more clear:
1. Add explicit error message for empty dictionary
    ```
        if isinstance(data, (tuple, list)):
            return find_batch_size(data[0])
        elif isinstance(data, Mapping):
            if len(data) == 0:
                raise ValueError(f"Can't find the batch size from empty {type(data)}.")
            for k in data.keys():
                return find_batch_size(data[k])
        elif not isinstance(data, torch.Tensor):
            raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
        return data.shape[0]
    ```
1. Use iterator instead of loop
    ```
        if isinstance(data, (tuple, list)):
            return find_batch_size(data[0])
        elif isinstance(data, Mapping):
            return next(iter(data.values()))  # <== will raise "StopIteration" on empty dict, which is a bit clearer than the original error
        elif not isinstance(data, torch.Tensor):
            raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
        return data.shape[0]
    ```
1. Combine the two above

What do you think on these improvements?


Discussion:
Thanks for highlighting this potentially confusing error. I think solution 1. looks cleanest to me.
Agree with solution 1 @YuryYakhno 

Resolving patch:
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
--- 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):


Test patch: