Add Flax example tests
# What does this PR do?

This PR adds tests for flax examples. 

The test for image-classification example is not added since the script uses a manual dataset with `torchvision`.  The test will be added once the script is refactored to use `datasets`.
