How use non-learnable parameters / freeze parameters in Flax/JAX? #1150
-
|
I.e. parameters that are part of the model, but not updated after initialization? (Original question by @dpkingma) |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 1 reply
-
|
The recommended solution is to use a Another approach is to use |
Beta Was this translation helpful? Give feedback.
-
|
|
Beta Was this translation helpful? Give feedback.
The recommended solution is to use a
MultiOptimizer, our documentation contains an example: https://flax.readthedocs.io/en/latest/flax.optim.html#multioptimizerAnother approach is to use
jax.lax.stop_gradient: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.stop_gradient.html