Fixes random_flax_module
with flax.linen.BatchNorm
#3394
Loading
random_flax_module
with flax.linen.BatchNorm
#3394