def METHOD_NAME(self, input):
    bw = self.weight
    ba = input
    bw = bw - bw.mean()
    self.scale = Parameter(
        (
            F.linear(ba, bw).std()
            / F.linear(torch.sign(ba), torch.sign(bw)).std()
        )
        .float()
        .to(ba.device)
    )
    # corner case when ba is all 0.0
    if torch.isnan(self.scale):
        self.scale = Parameter(
            (bw.std() / torch.sign(bw).std()).float().to(ba.device)
        )