Converting Pytorch GRU weights to TensorRT

Hello!I’m, trying to convert Pytorch weights to TensorRT weights for GRUCell.

I’m calling add_rnn_v2 with input shape [16, 1, 512], layer_count = 1 (as I just have one cell), hidden_size = 512, max_seq_len = 1, and op = trt.tensorrt.RNNOperation.GRU. It adds new layer successfully. After it, I call set_weights_for_gate for three gates: reset, update, hidden (as GRU only has 3 gates). Pytorch weights for GRUCell are in shape [1536, 512], dividing it in three I can get weights for every gate. So I call set_weights_for_gate with params: 0, trt.tensorrt.RNNGateType.UPDATE/RESET/HIDDEN, False, pytorch weight [512, 512]. The same for bias.
It doesn’t throw any error but when I call build_cuda_engine I get ERROR: Missing weights/bias: (0, RESET, 1). What am I doing wrong?

The mistake was that I called set_weights_for_gate only once for each gate. You need to call it one time for input weights (third parameter in this method) and one time for hidden.