There are assumptions in torch.nn.parallel.DistributedDataParallel today that prevent you from doing this unfortunately. We’re working on some changes to DDP to make this possible. Stay tuned.
cc @mrshenli