[Hot Fix] Ensure process_dataloader is called when tpu_cores > 1 to use Parallel DataLoader

[tchaton]

What does this PR do?

There was a regression in speed when using the TPU accelerator. After inspection, we realised that during the accelerator refactor we missed migrating the process_dataloader hook to the TrainingTypePlugin. This meant that we never wrapped the model when TPU cores > 1 in a parallel data-loader. This PR adds the correct logic.

A bigger discussion needs to take place in finding these corner cases which would be shown with regression/benchmarking tests. In the interim we should fix this particular error for release.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[Borda]

can we get a test for it?
also, edit chlog...?

[codecov]

Codecov Report

Merging #6015 (2401e3d) into master (fcfa7fa) will decrease coverage by 0%.
The diff coverage is 33%.

@@           Coverage Diff           @@
##           master   #6015    +/-   ##
=======================================
- Coverage      90%     90%    -0%     
=======================================
  Files         170     159    -11     
  Lines       11784   11209   -575     
=======================================
- Hits        10664   10140   -524     
+ Misses       1120    1069    -51     

[SeanNaren]

I made an edit to the main post, I think we need to have some benchmark tests for TPUs that ensure that we stay within reasonable limits for speed/accuracy. We test for model change, but would be good to have a general test for catching these regressions

[williamFalcon]

nice... yeah we need a test for this so we don't regress in the future