def replace_layer(model, layer_name, replace_fn):
    """Replace single layer in a (possibly nested) torch.nn.Module using `replace_fn`.
    Given a module `model` and a layer specified by `layer_name` replace the layer using
    `new_layer = replace_fn(old_layer)`. Here `layer_name` is a list of strings, each string
    indexing a level of the nested model."""
    if layer_name:
        nm = layer_name.pop()
        model._modules[nm] = replace_layer(model._modules[nm], layer_name, replace_fn)
    else:
        model = replace_fn(model)
    return model