keep reuqires_grad unchanged after converting bn to syncbn (#22569)

tianzhi0549 · facebook-github-bot · commit 6eb3969ac7f6 · 2019-07-10T08:38:04.000-07:00
Summary: After converting BN layers to SyncBN layers, the function will set all `requires_grad = True` regardless of the original requires_grad states. I think it is a bug and have fixed it in this PR. Pull Request resolved: #22569 Differential Revision: D16151647 Pulled By: zou3519 fbshipit-source-id: e2ad1886c94d8882485e7fb8be51ad76469ecc67
diff --git a/torch/nn/modules/batchnorm.py b/torch/nn/modules/batchnorm.py
@@ -493,6 +493,9 @@ def convert_sync_batchnorm(cls, module, process_group=None):
             if module.affine:
                 module_output.weight.data = module.weight.data.clone().detach()
                 module_output.bias.data = module.bias.data.clone().detach()
+                # keep reuqires_grad unchanged
+                module_output.weight.requires_grad = module.weight.requires_grad
+                module_output.bias.requires_grad = module.bias.requires_grad
             module_output.running_mean = module.running_mean
             module_output.running_var = module.running_var
             module_output.num_batches_tracked = module.num_batches_tracked
