Pytorch - Handling in place operation for sequence to sequence Multi-dimensional LSTM layer

I am trying to implement a sequence to sequence LSTM layer in Pytorch. Here a sequence can be of higher dimension than 1 (an image is a sequence, it just has 2 indexes that increase instead of one). More information in the following paper.

The forward() function of my network is:

def forward(self, x):
    """ Note: x is of shape (d1, ..., dn, batch_size, input_size). """
    dimensions = x.shape[:-2]
    batch_size = x.shape[-2]
    f = torch.empty(self.dim_in, *dimensions, batch_size, self.size_out)
    i = torch.empty(*dimensions, batch_size, self.size_out)
    o = torch.empty(*dimensions, batch_size, self.size_out)
    c = torch.empty(*dimensions, batch_size, self.size_out)
    s = torch.empty(*dimensions, batch_size, self.size_out)
    h = torch.empty(*dimensions, batch_size, self.size_out)
    for idx in self.iter_idx(dimensions):
        # 1/ Forget, input, output and cell activation gates.
        for l in range(self.dim_in):
            f[l][idx] = torch.sigmoid(self.biasf[l] + torch.mm(x[idx], self.wf[l]) + sum(torch.mul(h[prev(idx,k)], self.uf[l][k]) for k in np.nonzero(idx)[0]))
        i[idx] = torch.sigmoid(self.biasi + torch.mm(x[idx], self.wi) + sum(torch.mul(h[prev(idx,k)], self.ui[k]) for k in np.nonzero(idx)[0]))
        o[idx] = torch.sigmoid(self.biaso + torch.mm(x[idx], self.wo) + sum(torch.mul(h[prev(idx,k)], self.uo[k]) for k in np.nonzero(idx)[0]))
        c[idx] = torch.sigmoid(self.biasc + torch.mm(x[idx], self.wc) + sum(torch.mul(h[prev(idx,k)], self.uc[k]) for k in np.nonzero(idx)[0]))
        # 2/ cell state
        s[idx] = torch.tanh(torch.mul(i[idx], c[idx]) + sum(torch.mul(f[k][idx], s[prev(idx,k)]) for k in np.nonzero(idx)[0]))
        # 3/ Final output
        h[idx] = torch.mul(o[idx], s[idx])
    return h

When I try to run it and the backward() method of the loss is called, the following error is thrown:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [32, 3]], which is output 0 of SelectBackward, is at version 100; expected version 99 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

The backtrace points to this specific line:

s[idx] = torch.tanh(torch.mul(i[idx], c[idx]) + sum(torch.mul(f[k][idx], s[prev(idx,k)]) for k in np.nonzero(idx)[0]))

From what I understand, Pytorch is angry at me because I modified the tensor s in place. Looking at the error message, I see version 100 and expected version 99, which implies the error is thrown when computing the gradient for the last iteration of the loop and Pytorch expected to have s at a version ine loop iteration prior.

For me it shouldn't be an issue, because this loop merely "initializes" s (and the other tensors too). But I understand that the interpreter may have a hard time determining this statically.

So I have 3 questions:

• why s specifically? It is neither the first nor the last tensor to be modified in place during the loop.
• is there a way to tell Pytorch that I am just filling the arrays? Or to tell him to only look at the last version for gradient computation?
• is there a way to rewrite this code without in place operations? I tried a version where f, i, o, etc are lists that grow with the loop, but I get the same error on lines like s = torch.cat((s,s_tmp), axis=0).

Thank you in advance.