@Stefaanhess, depending on your network structure, your code might not work. For example, if the children of layer have children with batch norm children, you will never touch them.
I use this to recursively search for all batch norm children to update:
def recursive_batch_norm_update(child):
    
    if type(child) == torch.nn.BatchNorm1d:
        child.momentum = a * b**(self.epoch/c)
        return
    
    for children in child.children():
        lowest_child = recursive_layer(children)
    
    return

then just call
recursive_batch_norm_update(network)