This is how I did
def forward(self, x, lens):
        h0 = torch.randn(self.num_layers, x.size(0), self.hidden_size)

        emb = self.embedding(x)  
        packed = nn.utils.rnn.pack_padded_sequence(
            emb, lens, batch_first=True, enforce_sorted=False)

        packed, h0 = self.gru(packed, h0)
        out = nn.utils.rnn.pad_packed_sequence(packed, batch_first=True)[0]
            
        index = (lens-1).unsqueeze(-1).repeat(1,out.size(-1)).unsqueeze(1)
        out = torch.gather(out, 1, index).squeeze()
        out = self.lin(out)
        return out