Solved by ptrblck in post #6 
              
              
                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():
       …