Skip to content

model.patch_lstm_layers_for_onnx_export

Full path: schola.core.model.patch_lstm_layers_for_onnx_export

Attach forward hooks so ONNX-exported LSTM hidden states match PyTorch.

patch_lstm_layers_for_onnx_export(module)

Parameters

  • module (torch.nn.Module) - Root module; every nested torch.nn.LSTM receives reshape_lstm_output_hook.