Function replace_submodules

Replace all submodules selected by the predicate with the output of func. predicate: Return true if the module is to be replaced. func: Return new module to use.

(
    root_module: nn.Module,
    predicate: Callable[[nn.Module], bool],
    func: Callable[[nn.Module], nn.Module],
)