Jax example
CNN #
Bases: Module
A simple CNN model.
Taken from https://flax.readthedocs.io/en/latest/quick_start.html#define-network
Bases: Module
A simple CNN model.
Taken from https://flax.readthedocs.io/en/latest/quick_start.html#define-network