RNN weights error with TensorRT4

Hi, I create 2 bi-GRU (hidden_size=512) layers network with network.add_rnnv2 in python, and load TF RNN weights to the network. The input of rnn layer is [1, 126, 768]. So for the first layer, each gate weight is [768512+512512]. Then for the second layer, each gate weight is [1024512+512512].

But after i set gate weights and biases, gru.get_output(0).get_dimensions() comes to an error:
[TensorRT] ERROR: (Unnamed Layer* 7) [RNN]: (1, RESET, 1) weights weights has count 524288 but 393216 was expected
(1, RESET, 1) means the reset gate weights in the second GRU layer, and its size is 1024512=524288, but it shows 393216(768512) was expected, which is the size of reset gate weights in the first GRU layer.
I am not sure if it is a bug or something wrong with API implement.

Below is the related codes:

shuf1.set_reshape_dimensions(trt.infer.Dims3(-1, src_length, rnn_inputdim))
gru = network.add_rnnv2(shuf1.get_output(0), 2, hidden_size, src_length, trt.infer.RNNOperation_kGRU)
gru.set_direction(trt.infer.RNNDirection_kBIDIRECTION)

for i in range(4):
layer = i // 2
if layer == 0:
inputdim = rnn_inputdim #first layer input
else:
inputdim = hidden_size * 2 #BIRNN
rows = inputdim + hidden_size
cols = hidden_size + hidden_size
tmpweight = weight_gate.reshape(rows, cols)
W_z = tmpweight[:inputdim, :hidden_size].transpose().reshape(-1)
R_z = tmpweight[inputdim:, :hidden_size].transpose().reshape(-1)
W_r = tmpweight[:inputdim, hidden_size:].transpose().reshape(-1)
R_r = tmpweight[inputdim:, hidden_size:].transpose().reshape(-1)
W_h = weight_input.reshape(inputdim, hidden_size).transpose().reshape(-1)
R_h = weight_hidden.reshape(hidden_size, hidden_size).transpose().reshape(-1)
b_zW = bias_gate[:hidden_size]
b_zR = np.zeros(hidden_size, dtype=np.float32)
b_rW = bias_gate[hidden_size:]
b_rR = np.zeros(hidden_size, dtype=np.float32)
b_hW = bias_input
b_hR = bias_hidden

    gru.set_weights_for_gate(layer, trt.infer.RNNGateType_kUPDATE, True, trt.infer.Weights(W_z))
    [b]gru.set_weights_for_gate(layer, trt.infer.RNNGateType_kRESET, True, trt.infer.Weights(W_r))[/b]
    gru.set_weights_for_gate(layer, trt.infer.RNNGateType_kHIDDEN, True, trt.infer.Weights(W_h))
    gru.set_bias_for_gate(layer, trt.infer.RNNGateType_kUPDATE, True, trt.infer.Weights(b_zW))
    gru.set_bias_for_gate(layer, trt.infer.RNNGateType_kRESET, True, trt.infer.Weights(b_rW))
    gru.set_bias_for_gate(layer, trt.infer.RNNGateType_kHIDDEN, True, trt.infer.Weights(b_hW))

    gru.set_weights_for_gate(layer, trt.infer.RNNGateType_kUPDATE, False, trt.infer.Weights(R_z))
    gru.set_weights_for_gate(layer, trt.infer.RNNGateType_kRESET, False, trt.infer.Weights(R_r))
    gru.set_weights_for_gate(layer, trt.infer.RNNGateType_kHIDDEN, False, trt.infer.Weights(R_h))
    gru.set_bias_for_gate(layer, trt.infer.RNNGateType_kUPDATE, False, trt.infer.Weights(b_zR))
    gru.set_bias_for_gate(layer, trt.infer.RNNGateType_kRESET, False, trt.infer.Weights(b_rR))
    gru.set_bias_for_gate(layer, trt.infer.RNNGateType_kHIDDEN, False, trt.infer.Weights(b_hR))

dims2 = gru.get_output(0).get_dimensions() #error!
[TensorRT] ERROR: (Unnamed Layer* 7) [RNN]: (1, RESET, 1) weights weights has count 524288 but 393216 was expected