Make PT/Flax tests could be run on GPU
# What does this PR do?

We don't have jax/flax on our CI runner, so no issue. But when trying to run those tests on a GPU machine with jax/flax installed, we have error. See comment.
_The documentation is not available anymore as the PR was closed or merged._