Honestly, I haven’t tried this code, just found it when researching to answer your question. It’s possible that there is a bug in datasets?
We are going to add this feature to all HF Trainer-based examples [examples] [distributed] process datasets.map only on main process in · Issue #12345 · huggingface/transformers · GitHub 6 so watch that Issue, I will try to work on it shortly and see what needs to be done to make it work.