nn.DataParallel would update the batchnorm stats on the default device, as seen in this code snippet:

bn = nn.BatchNorm2d(3, momentum=1.0).cuda()

a = torch.randn(16, 3, 224, 224) * 12 + 7
b = torch.randn(16, 3, 224, 224) * 5 - 9
x = torch.cat((a, b), dim=0)

model = nn.DataParallel(bn, device_ids=[0, 1])
out = model(x)
print(model.module.running_mean)
> tensor([6.9984, 7.0203, 7.0044], device='cuda:0')
print(torch.sqrt(model.module.running_var))
> tensor([11.9902, 12.0097, 11.9818], device='cuda:0')

If you are using nn.DistributedDataParallel (which would also be faster than nn.DataParallel), you could use SyncBatchNorm to synchronize these running stats.

You are directly passing the module to nn.DataParallel, which should be executed on multiple devices. E.g. if you only want to pass a submodule to it, you could use:

model = MyModel()
model.submodule = nn.DataParallel(model.submodule)


Transferring the parameters to the device after the nn.DataParallel creation should also work. I think the documentation mentions, that you would have to transfer the model to the default device at one point before executing the forward pass.