flax

484cc116 - Merge pull request #301 from ayush-1506:batchnorm_error_msg