Hi, I am trying to use GRU to predict tokens. However, the computation speed is very slow in for loop.
My code is as below:
decoded_words, word_distributions = [], []

for _ in range(max_utter_len):
    word_distribution, decoder_h = self.decoder(
        decoder_inputs, contexts, decoder_h)

    word_distribution = word_distribution.squeeze(1)
    # [batch size, voc size]

    top1 = word_distribution.topk(1)[1]
    # [batch size, 1]

    decoded_words.append(top1.squeeze(-1))
    word_distributions.append(word_distribution)

    decoder_inputs = top1

decoded_words = torch.stack(decoded_words, dim=-1)
word_distributions = torch.stack(word_distributions, dim=1)

It takes 17 seconds for every 100 iterations.
However, I use the “teacher forcing” in the training, and it only takes 7 seconds for every 100 iterations.
The code is as below:
word_distributions, decoder_init_hidden = self.decoder(
            inputs, context, decoder_init_hidden)

How to speed up RNN prediction?
Thank you.