Helper: Construct a model from hyperparameters.

def make_model(num_layers: int = 6,
               input_size: int = 512,  # Attention size
               hidden_size: int = 2048,  # FF layer size
               heads: int = 8,
               dropout: float = 0.1,
               return_all_layers: bool = False) -> TransformerEncoder:
    """Helper: Construct a model from hyperparameters."""
    attn = MultiHeadedAttention(heads, input_size, dropout)
    ff = PositionwiseFeedForward(input_size, hidden_size, dropout)
    model = TransformerEncoder(EncoderLayer(input_size, attn, ff, dropout),
                               num_layers,
                               return_all_layers=return_all_layers)

    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            torch.nn.init.xavier_uniform_(p)
    return model