thanks, this solved my problem. 
def forward(self, x):
    if not self.initialized:
        mean_x = torch.mean(x)
        torch.distributed.all_reduce(mean_x, torch.distributed.ReduceOp.SUM)
        a = nn.Parameter(mean_x / torch.distributed.get_world_size())
        self.initialized = True