Pytorch - how to wrap LSTM layer?

I have a problem with applying some XAI-algorithms of Captum on my LSTM. Got the following advice:

'what about wrapping self.lstm with a layer that only returns a out instead of out, (h_out, c_out)?'

Unfortunately I do not have a programming background and do not really know how to implement this advice.

Someone may be able to help?

class LSTM(nn.Module):

    def __init__(self, input_size, hidden_size, num_layers, output_size = 1,dropout_prob = 0.2):
        super(LSTM, self).__init__()
        
        self.input_size = input_size  #number of features used for input (e.g. 3: Precip, T, humid)
        self.hidden_size = hidden_size #number of neurons per layer
        self.num_layers = num_layers  #number of layer
        self.dropout = dropout_prob #dropout rate against overfitting

        
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size,
                            num_layers=num_layers, batch_first=True)
        
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        #set initial values for cell/hidden state
        h_0 = Variable(torch.zeros(
            self.num_layers, 1, self.hidden_size))
        
        c_0 = Variable(torch.zeros(
            self.num_layers, 1, self.hidden_size))
       

        # Propagate input through LSTM
        out,(h_out,c_out) = self.lstm(x)

        out = out[ :,-1, :]
        out = self.fc(out)
        return out

Thanks a lot!