def xavier_init(module, gain=1, bias=0, distribution='normal'):
    assert (distribution in ['uniform', 'normal'])
    if (hasattr(module, 'weight') and (module.weight is not None)):
        if (distribution == 'uniform'):
            nn.init.xavier_uniform_(module.weight, gain=gain)
        else:
            nn.init.xavier_normal_(module.weight, gain=gain)
    if (hasattr(module, 'bias') and (module.bias is not None)):
        nn.init.constant_(module.bias, bias)
def xavier_init(module, gain=1, bias=0, distribution='normal'):
    assert (distribution in ['uniform', 'normal'])
    if (hasattr(module, 'weight') and (module.weight is not None)):
        if (distribution == 'uniform'):
            nn.init.xavier_uniform_(module.weight, gain=gain)
        else:
            nn.init.xavier_normal_(module.weight, gain=gain)
    if (hasattr(module, 'bias') and (module.bias is not None)):
        nn.init.constant_(module.bias, bias)