What is the proper way to use BnFreeze callback

Hi, I am training on large images of size 1024*1024 for an image segmentation problem. Since I am able to use a batch size of 4 on 4gpu’s, I wanted to freeze the Batch normalization layer. When looking at the source code of BnFreeze callback I realized the function expects the gradients of the BN layer to be false. I could not find a function which does that in fastai. I remember using it in the previous version of fastai. So I wrote a function which does that. Wanted to check if I am doing it in a proper/recommended way.

Function to freeze BN layers.

def freeze_bn(m):
    for l in children(m):
        if isinstance(l, bn_types):
            for o in l.parameters():
                o.requires_grad = False

I pass model to freeze_bn function and then use BnFreeze callback.