This dummy code snippet seems to work, if I use updateBN and BN_grad_zero after the backward call:
def updateBN():
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.weight.grad.add_(torch.sign(m.weight)) 

def BN_grad_zero():
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            mask = (m.weight != 0)
            mask = mask.float()
            m.weight.grad.mul_(mask)
            m.bias.grad.mul_(mask)
            

model = nn.Sequential(
    nn.Conv2d(1, 3, 3, 1, 1),
    nn.BatchNorm2d(3))

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
optimizer.zero_grad()
output = model(torch.randn(1, 1, 24, 24))
loss = F.cross_entropy(output, torch.randint(0, 3, (1, 24, 24)))
pred = output.data.max(1, keepdim=True)[1]
loss.backward()

updateBN()
BN_grad_zero()
optimizer.step()

PS: I’ve also removed the .data attribute, as we don’t recommend to use it, since it might yield unwanted side effects.