Hi, there, I am working on a new RNN unit implementation. Since the formulation is totally different with existing RNN units, I implemented everything from scratch. In order to process information in each time stamp, I used a for loop to loop through time stamps. It looks like the codes below. Unfortunately, it is much slower then its theano counterpart. I am wondering is there a special way to do the for loop in RNNs which can be faster than the naive way. So, in a more general case, is there any special way to speed up general for loop process in pytorch?
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(RNN, self).__init__()
        self.hidden_size = hidden_size
        
        self.i2h = nn.Linear(input_size+hidden_size, hidden_size)
        self.h2o = nn.Linear(input_size+hidden_size, output_size)
        self.tanh = nn.Tanh()

    def forward(self, X):
        time_steps = X.size(0)
        batch_size = X.size(1)
        hidden = Variable(torch.zeros(batch_size, self.hidden_size))
        outputs = []
        hiddens = []
        for t in range(time_steps):
            x_input = X[t]
            hidden_input = hidden
            inp = torch.cat( (x_input,hidden_input), 1 )
            hidden = self.tanh(self.i2h(inp))
            output = self.h2o(inp)
            outputs.append(output)
            hidden.append(hidden)
        return torch.cat(hiddens, 1), torch.cat(outputs, 1)