Thanks a lot. I have one last question. Like the OP, I need to recreate the state dict every time in the forward pass. I see about 8x increase in training time when compared to original PyTorch DataParallel. Any ideas why this might be the case?
def create_state_dict_new(main_module):
    state_dict_data = OrderedDict()

    def state_dict_recursion(this_module, state_dict_data, prefix=''):
        if hasattr(this_module,"_former_parameters"):
            for name, param in this_module._former_parameters.items():
                if param is not None:
                    state_dict_data[prefix + name] = param
        for name, buf in this_module._buffers.items():
            if buf is not None:
                state_dict_data[prefix + name] = buf
        for name, module in this_module._modules.items():
            if module is not None:
                state_dict_recursion(module, state_dict_data, prefix + name + '.')
    state_dict_recursion(main_module._modules['model'], state_dict_data)
    return state_dict_data

class ModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super(ModelWrapper, self).__init__()
        self.model = model
    def forward(self, x):
        state_list = create_state_dict_new(self)
        return model(x)

model = torch.nn.DataParallel(ModelWrapper(model))