openfold.model.torchscript¶
Functions
|
TorchScript a handful of low-level but frequently used submodule types that are known to be scriptable. |
|
Convert all submodules whose types match one of those in the input list to recursively scripted equivalents in place. |
- script_preset_(model)¶
TorchScript a handful of low-level but frequently used submodule types that are known to be scriptable.
- Parameters:
model (Module) – A torch.nn.Module. It should contain at least some modules from this repository, or this function won’t do anything.
- script_submodules_(model, types=None, attempt_trace=True, batch_dims=None)¶
Convert all submodules whose types match one of those in the input list to recursively scripted equivalents in place. To script the entire model, just call torch.jit.script on it directly.
When types is None, all submodules are scripted.
- Parameters:
model (Module) – A torch.nn.Module
types (Sequence[type] | None) – A list of types of submodules to script
attempt_trace (bool | None) – Whether to attempt to trace specified modules if scripting fails. Recall that tracing eliminates all conditional logic—with great tracing comes the mild responsibility of having to remember to ensure that the modules in question perform the same computations no matter what.