Pytorch, slicing tensor causes RuntimeError:: one of the variables needed for gradient computation has been modified by an inplace operation:

I wrote a RNN with LSTM cell with Pycharm. The peculiarity of this network is that the output of the RNN is fed into a integration opeartion, computed with Runge-kutta. The integration takes some input and propagate that in time one step ahead. In order to do so I need to slice the feature tensor X along the batch dimension, and pass this to the Runge-kutta.

class MyLSTM(torch.nn.Module):
    def __init__(self, ni, no, sampling_interval, nh=10, nlayers=1):
        super(MyLSTM, self).__init__()

        self.device = torch.device("cpu")
        self.dtype = torch.float
        self.ni = ni
        self.no = no
        self.nh = nh
        self.nlayers = nlayers

        self.lstms = torch.nn.ModuleList(
            [torch.nn.LSTMCell(self.ni, self.nh)] + [torch.nn.LSTMCell(self.nh, self.nh) for i in range(nlayers - 1)])
        self.out = torch.nn.Linear(self.nh, self.no)
        self.do = torch.nn.Dropout(p=0.2)
        self.actfn = torch.nn.Sigmoid()
        self.sampling_interval = sampling_interval
        self.scaler_states = None
        # Options

    # description of the whole block
    def forward(self, x, h0, train=False, integrate_ode=True):
        x0 = x.clone().requires_grad_(True)
        hs = x  # initiate hidden state

        if h0 is None:
            h = torch.zeros(hs.shape[0], self.nh, device=self.device)
            c = torch.zeros(hs.shape[0], self.nh, device=self.device)
        else:
            (h, c) = h0

        # LSTM cells
        for i in range(self.nlayers):
            h, c = self.lstms[i](hs, (h, c))
            if train:
                hs = self.do(h)
            else:
                hs = h

        # Output layer
        # y = self.actfn(self.out(hs))
        y = self.out(hs)

        if integrate_ode:
            p = y
            y = self.integrate(x0, p)
        return y, (h, c)

    def integrate(self, x0, p):
        # RK4 steps per interval
        M = 4
        DT = self.sampling_interval / M
        X = x0
        # X = self.scaler_features.inverse_transform(x0)

            for b in range(X.shape[0]):
                xx = X[b, :]
                for j in range(M):
                    k1 = self.ode(xx, p[b, :])
                    k2 = self.ode(xx + DT / 2 * k1, p[b, :])
                    k3 = self.ode(xx + DT / 2 * k2, p[b, :])
                    k4 = self.ode(xx + DT * k3, p[b, :])
                    xx = xx + DT / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
                X_all[b, :] = xx
        return X_all

    def ode(self, x0, y):
        # Here I a dynamic model

I get this error:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor []], which is output 0 of SelectBackward, is at version 64; expected version 63 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

the problem is in the operations xx = X[b, :] and p[b,:]. I know that because I choose batch dimension of 1, then I can replace the previous two equations with xx=X and p, and this works. How can split the tensor without loosing the gradient?