towards_data_science_44.py

import padl
import torch


hidden_size = padl.param('hidden_size', 512)
input_size = padl.param('input_size', 64)
n_tokens = padl.param('n_tokens', 16)
nn = padl.transform(torch.nn)


@padl.transform
class HiddenState(torch.nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def forward(self, x):
        return self.layer(x)[0]


rnn_layer = nn.GRU(input_size, hidden_size, 1)
_pd_main = (
    nn.Embedding(n_tokens, input_size)
    >> HiddenState(rnn_layer)
    >> nn.Linear(hidden_size, 1)
    >> nn.Sigmoid()
)