Optional type of lengths causes slow speed in LengthGroupedSampler
Author: seopbo
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.20.1
- Platform: macOS-12.4-arm64-arm-64bit
- Python version: 3.8.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: no

but this bug doesn't depend on environment.

### Who can help?

@sgugger 

### Information

- [x] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python
import time
import torch
import random

megabatches = [[random.randint(1, 512) for _ in range(51200)] for _ in range(2)]

# test lengths_list_int
lengths_torch_tensor = torch.tensor([random.randint(1, 512) for _ in range(102400)])

start = time.time()
megabatches_ = [list(sorted(megabatch, key=lambda i: lengths_torch_tensor[i], reverse=True)) for megabatch in megabatches]
end = time.time()
print(end-start)

# test lengths_list_int
lengths_torch_tensor = torch.tensor([random.randint(1, 512) for _ in range(102400)])
lengths_list_int = lengths_torch_tensor.tolist()

start = time.time()
megabatches_ = [list(sorted(megabatch, key=lambda i: lengths_list_int[i], reverse=True)) for megabatch in megabatches]
end = time.time()
print(end - start)
```
```bash
1.0904111862182617
0.013269901275634766
```

### Expected behavior

When using `group_by_length`, `length_column_name` in `TrainingArguments`, `get_length_grouped_indices` function in `LengthGroupedSampler` is very slow if `Dataset[length_cloumn_name]` is `torch.Tensor(List[int])` (e.g. `torch.Tensor([200,100,..])`). So I think that `lengths: Optional[List[int]]` -> `lengths: List[int]` in `__init__` method of `LengthGroupedSampler` or warning message is printed and type casting to `List[int]`.

https://github.com/huggingface/transformers/blob/49c8c67fb815a277405f84dea4a66353e19fb347/src/transformers/trainer_pt_utils.py#L532-L569

https://github.com/huggingface/transformers/blob/49c8c67fb815a277405f84dea4a66353e19fb347/src/transformers/trainer_pt_utils.py#L520

Comments:
I don't understand why changing the type annotation would change anything. Those type annotations are completely ignored at execution, they are just here as documentation of the code. Python is not a typed language.

The problem is just that if a tensor is passed for the lengths, it should be converted back to a list. Would you like to make a PR with that?


> I don't understand why changing the type annotation would change anything. Those type annotations are completely ignored at execution, they are just here as documentation of the code. Python is not a typed language.
> 
> The problem is just that if a tensor is passed for the lengths, it should be converted back to a list. Would you like to make a PR with that?

Ok, sure. I'll pr this. @sgugger