Added torch.no_grad() to update_bn (#52654)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/52055
This fixes the **out of memory error** while using update_bn in **SWA**, by not allocating memory for backpropagation.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52654
Reviewed By: malfet
Differential Revision: D26620077
Pulled By: albanD
fbshipit-source-id: 890b5a78ba9c1a148f3ab7c63472a73d8f6412a4