Solved by ptrblck in post #2 
              
              
                If I’m not mistaken, the recursive loop would enable the gradient computation for all nested BatchNorm2d layers even in the custom modules, as they would be reached last. 
You could verify it manually after using this approach via e.g.: 
print(model.custom_layer.bn_layer.weight.requires_grad)