flax
Use optax's losses and schedules in mnist and imagenet flax examples.
#1286
Merged
Merged