[Flax examples] remove dependancy on pytorch training args
# What does this PR do?

Remove dependency on PyTorch `TrainingArguments` class in flax examples.

Fixes #13721
@patrickvonplaten I've run the flax examples tests and made sure that they pass with this PR. Merging this now.
Feel free to leave comments if you see something strange :) <|||||>Awesome - thanks @patil-suraj 