Fix barrier for SM distributed
# What does this PR do?

#12351 introduced a new context manager for having the main process execute an instruction while other process have to wait. That context manager was missing special treatment for TPUs (added in #12464) and SageMaker distributed. This PR adds the latter.

Fixes #12847