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

[franz101]

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, ...)
• 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

[franz101]

Reproducible in this colab example currently

[franz101]

Found the issue, it seems like the on step end is called only after two data points have been collated.