Fix gather for TPU
# What does this PR do?

This PR fixes #13798, which is cause by the gather of trainer loss, a 0d-tensor. The code does not add one dimension to those to concatenate properly, resulting in the error.
Thanks for fixing it - sorry should have tested on TPU as well :-/<|||||>No worries!