espnet2.torch_utils.compile_modules.compile_modules

source

espnet2.torch_utils.compile_modules.compile_modules(model: Module, names: Sequence[str], conf: Dict | None = None)

Compile the named submodules of model in place with torch.compile.

Each name is a dotted path below model, e.g. “encoder” or “frontend.logmel”. This function calls nn.Module.compile() on each one, which compiles the module in place: the parameter names and the state dict do not change, so a checkpoint saved from the compiled model loads into an uncompiled one. Compiling the modules that do the work, and not the whole model, leaves the loss and the statistics in eager mode.

• Parameters:
  • model – The model.
  • names – Dotted paths of the submodules to compile.
  • conf – Keyword arguments of torch.compile, e.g. {“mode”: “max-autotune”}.
• Raises:ValueError – A name is empty (the whole model), is not a submodule of model, is given twice, or lies below another one (compiling a module and a module inside it would compile the inner one twice).