def METHOD_NAME(ntokens,
                 embedding_dim=200,
                 memory_size=128,
                 num_layers=4,
                 num_heads=4):
    embedding_layer = torch.nn.Embedding(ntokens, embedding_dim)
    embedding_layer.weight.data.uniform_(-0.1, 0.1)
    return networks.Sequential(
        embedding_layer,
        layers.Reshape(1, -1),
        networks.TransformerNetwork(
            input_tensor_spec=alf.TensorSpec((1, embedding_dim)),
            num_prememory_layers=0,
            num_attention_heads=num_heads,
            d_ff=4 * embedding_dim,
            core_size=1,
            use_core_embedding=False,
            memory_size=memory_size,
            num_memory_layers=num_layers,
            centralized_memory=False),
        layers.FC(embedding_dim, ntokens, kernel_init_gain=0.),
        input_tensor_spec=alf.TensorSpec((), dtype=torch.int64),
    )