How can I change the model form LSTM to GRU?

[Pluto-yiyi]

Hi, I have run this code, and test successfully. But when I want to change the model to GRU, I can't finish it. I hope you can help me and give some suggests. Thank you very much!
I change the code like this:
【train.py】
before: parser.add_argument('--model', type=str, default='lstm',
help='lstm, rnn, gru, or nas')
after:parser.add_argument('--model', type=str, default='gru',
help='lstm, rnn, gru, or nas')