#24028 seems to break the last epoch for a logging integration
### System Info

Hey @muellerzr,

thanks for your lightning fast (accelerated) reply ;) regarding #24028, I'm currently debugging what's causing the issue

Setup:
- A custom callback to log embeddings, the data collator in the Trainer is wrapped to extract ids of each sample in a batch

Error:
- The wrapped data collation works fine except in the last step

How to reproduce?
See reproduction tab

Currently this is the example I can show for reproduction.
My first guess, it's related to multiprocessing. It seems like the custom collator is not called in the last step.
But will give more details or a possible solution soon.


### Who can help?

@muellerzr

### Information

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

### Tasks

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

### Reproduction

```
git clone https://github.com/rungalileo/dataquality.git
cd dataquality
python -m venv .venv
source .venv/bin/activate
pip install invoke
inv all
pip install --upgrade transformers
pytest tests/integrations/hf/test_text_classification_hf.py -s -k test_remove_unused_columns
```

### Expected behavior

Test should finish collating each step
Reproducible in [this colab example](https://colab.research.google.com/drive/1NObtjxY37VEx2u3SBFDDa_zwTILbX3B2?usp=sharing) currently<|||||>Found the issue, it seems like the on step end is called only after two data points have been collated.