Skip to content

model.reshape_lstm_output_hook

Full path: schola.core.model.reshape_lstm_output_hook

Reshape LSTM hidden states during ONNX export so hn / cn match PyTorch layouts.

reshape_lstm_output_hook(lstm, args, output)

Parameters

  • lstm (torch.nn.LSTM) - Layer instance receiving the hook.

  • args (tuple) - Forward inputs (input, (h_0, c_0)) as passed to LSTM.forward.

  • output (tuple) - Forward outputs (output, (h_n, c_n)).