if module.bias is not None:
    if len(out.shape) == 4:
        out = out + module.bias.view(1, -1, 1, 1)
    if len(out.shape) == 3:
        out = out + module.bias.view(1, -1, 1)
    if len(out.shape) == 2:
        out = out + module.bias.view(1, -1)