-record(bert_embedding, {
    word_embeddings :: viva_tensor@nn@embedding:embedding(),
    position_embeddings :: viva_tensor@nn@embedding:learned_positional_encoding(),
    token_type_embeddings :: viva_tensor@nn@embedding:embedding(),
    layer_norm :: viva_tensor@nn@norm:layer_norm()
}).