Retrieve only the last hidden state from lstm layer in pytorch sequential
I have a pytorch model:
model = torch.nn.Sequential(
        torch.nn.LSTM(40, 256, 3, batch_first=True),
        torch.nn.Linear(256, 256),
        torch.nn.ReLU()
)

And for the LSTM layer, I want to retrieve only the last hidden state from the batch to pass through the rest of the layers. Ex:
_, (hidden, _) = lstm(data)
hidden = hidden[-1]

Though, that example only works for a subclassed model. I need to somehow do this on a nn.Sequential() model that way when I save it, it can properly be converted to a tensorflow.js model. The reason I can't make and train this model in tensorflow.js is because I'm trying to implement this repo: Resemblyzer in tensorflow.js while still using the same weights as the pretrained Resemblyzer model which was made in pytorch as a subclassed model. I thought of using the torchvisions.transformations.Lambda() transformation but I would assume that would make it incompatible with tensorflow.js. Is there any way to make this possible while still allowing the model to convert properly?

You could split up your sequential but only doing so in the forward definition of your model on inference. Once defined:
model = nn.Sequential(nn.LSTM(40, 256, 3, batch_first=True),
                      nn.Linear(256, 256),
                      nn.ReLU())

You can split it:
&gt;&gt;&gt; lstm, fc = model[0], model[1:]

Then infer in two steps:
&gt;&gt;&gt; out, (hidden, _) = lstm(data)
&gt;&gt;&gt; hidden = hidden[-1]
&gt;&gt;&gt; out = fc(out)       # &lt;- or fc(out[-1]) depending on what you want

