[Reland] Mixed precision batchnorm fix (#77234)
Reland of https://github.com/pytorch/pytorch/pull/77089, which was reverted due to land race.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77234
Approved by: https://github.com/zhaojuanmao
