Export RNN (GRU) as ONNX, how to set the batch size

1

Hello to all!
I’m trying to export a trained GRU as ONNX. As I want to test it online (take each timestep in a loop and hand over hidden state), I would like to have a fixed input length of 1. With batch_first=True I use an input tensor of size (batch_sz=1, seq_len=1, feat=10). When I try to export the ONNX I get the Warning:

"/Users//miniforge3/envs/arm_native/lib/python3.9/site-packages/torch/onnx/symbolic_opset9.py:1909: UserWarning: Exporting a model to ONNX with a batch_size other than 1, with a variable length with GRU can cause an error when running the ONNX model with a different batch size. Make sure to save the model with a batch size of 1, or define the initial states (h0/c0) as inputs of the model.
warnings.warn("Exporting a model to ONNX with a batch_size other than 1, " + "

What am I doing wrong? I also tried to set dynamic axis for the input-length, with no effect.
I can’t find any examples for exporting recurrent nets, only old workarounds for when torch didn’t support this feature.

import torch

class Test(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.GRU(input_size=10, hidden_size=5, num_layers=2, batch_first=True)
        self.net.eval()

    def forward(self, x, h):
        return self.net(x, h)


m = Test()
_x = torch.rand((1, 1, 10))
_h = torch.rand((2, 1, 5))
_out, _h_n = m(_x, _h)

torch.onnx.export(m,
                  (_x, _h),
                  'def.onnx',
                  export_params=True)