How do I enforce bfloat16 activations in my model? #4723
-
With the new NNX API it's not super clear to me how I can enforce that the activations of my nnx model are in another dtype than the models weights. Are there best practices for mixed precision training? |
Beta Was this translation helpful? Give feedback.
Answered by
wittenator
May 13, 2025
Replies: 1 comment
-
Just for reference: I missed that there are dtype and param_dtype parameters for each layer. Dtype specifies the activation dtype and param_dtype the dtype of the weights. |
Beta Was this translation helpful? Give feedback.
0 replies
Answer selected by
wittenator
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Just for reference: I missed that there are dtype and param_dtype parameters for each layer. Dtype specifies the activation dtype and param_dtype the dtype of the weights.