Remove loss from some flax models docs & examples
# What does this PR do?

Tiny change: remove `loss` & `return_loss` in flax docs & examples.

@patil-suraj 
_The documentation is not available anymore as the PR was closed or merged._