Guard distributed imports (#77727)
Move distributed import after dist.is_avail check to fix builds with USE_DISTRIBUTED=0. Although, note that this issue is not caught by any CI at the moment.
Closes https://github.com/pytorch/pytorch/issues/77704
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77727
Approved by: https://github.com/malfet