Stop calling expand_1d on newer TF versions
Tensorflow changed their default `train_step` in version 2.11 to no longer user `data_adapter.expand_1d`, and also deleted that method. Since we copied that code for our train step, this made our `train_step` stop working in 2.11 when the user was using a non-dummy loss!

This PR resolves the issue by not calling `expand_1d` for TF versions >= 2.11.

Fixes #20750 
_The documentation is not available anymore as the PR was closed or merged._