How to freeze embedding layers

[lintangsutawika]

Hi,

I'm wondering how I might freeze token embedding layers in Unitransformer implementations. All references online seem to point to keras and not implementations with mesh.

https://github.com/tensorflow/mesh/blob/52a2332c3bb0aa5898caba7efecc8cfa0486276e/mesh_tensorflow/transformer/transformer.py#L697

Thank you