How to count the parameters in a Flax model? #1224

Answered by marcvanzee

marcvanzee asked this question in Q&A

marcvanzee
Apr 8, 2021
Maintainer

Original question by @ameya98.

Answered by marcvanzee

Answer by @jheek:

sum(p.size for p in jax.tree_leaves(params))

View full answer

Replies: 1 comment

marcvanzee
Apr 8, 2021
Maintainer Author

Answer by @jheek:

sum(p.size for p in jax.tree_leaves(params))

0 replies

Answer selected by marcvanzee

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment