tensorflow / mesh

Mesh TensorFlow: Model Parallelism Made Easier
Apache License 2.0
1.58k stars 254 forks source link

How to freeze embedding layers #364

Open lintangsutawika opened 3 years ago

lintangsutawika commented 3 years ago

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