towards_data_science_38.py

import torch

        
class HiddenState(torch.nn.Layer):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer
        
    def forward(self, x):
        return self.layer(x)[0]


def build_classifier(
    rnn_layer,
    n_tokens,
):

    return torch.nn.Sequential(
        torch.nn.Embedding(n_tokens, n_embed),
        HiddenState(rnn_layer),
        torch.nn.ReLU(),
        torch.nn.Linear(rnn_layer.hidden_size, 1),
        torch.nn.Sigmoid(),
    )