Save hidden states just before prediction

Hello, I’m a newbie to pytorch, I’m trying to do the following :

  • Create a function that takes as an input my text and returns a list of the indices of my tokens.
  • Create a function that has as an input batch size ( batch_size, seq_len) and returns hidden states : (batch_size, hidden_states)
    I included my model and the forward function that I belive have to make a forward_prime function that stops right before the prediction layer and save the hidden states , but I don’t know how to do it.
class ClassificationMultiLabel(nn.Module):
    def __init__(self, op_size, n_tokens, pretrained_vectors, nl=2, bidirectional=True, emb_sz=300, n_hiddenUnits=100):
        super(ClassificationMultiLabel, self).__init__()
        self.n_hidden = n_hiddenUnits
        self.embeddings = nn.Embedding(n_tokens, emb_sz)
        #         self.embeddings.weight.requires_grad = False
        self.rnn = nn.LSTM(emb_sz, n_hiddenUnits, num_layers=2, bidirectional=True, dropout=0.2)
        self.lArr = []
        if bidirectional:
            n_hiddenUnits = 2 * n_hiddenUnits
        self.bn1 = nn.BatchNorm1d(num_features=n_hiddenUnits)
        for i in range(nl):
            if i == 0:
                self.lArr.append(nn.Linear(n_hiddenUnits * 3, n_hiddenUnits))
                self.lArr.append(nn.Linear(n_hiddenUnits, n_hiddenUnits))
        self.lArr = nn.ModuleList(self.lArr)
        self.l1 = nn.Linear(n_hiddenUnits, op_size)

    def forward(self, data, lengths):
        bs = data.shape[1]
        self.h = self.init_hidden(bs)
        embedded = self.embeddings(data)
        embedded = nn.Dropout()(embedded)
        #         embedded = pack_padded_sequence(embedded, torch.as_tensor(lengths))
        rnn_out, self.h = self.rnn(embedded, (self.h, self.h))
        #         rnn_out, lengths = pad_packed_sequence(rnn_out,padding_value=1)
        avg_pool = F.adaptive_avg_pool1d(rnn_out.permute(1, 2, 0), 1).view(bs, -1)
        max_pool = F.adaptive_max_pool1d(rnn_out.permute(1, 2, 0), 1).view(bs, -1)
        ipForLinearLayer =[avg_pool, max_pool, rnn_out[-1]], dim=1)
        for linearlayer in self.lArr:
            outp = linearlayer(ipForLinearLayer)
            ipForLinearLayer = self.bn1(F.relu(outp))
            ipForLinearLayer = nn.Dropout(p=0.6)(ipForLinearLayer)
        outp = self.l1(ipForLinearLayer)
        return outp

full code can be found here :