flax

f34030e5 - Adds a comment on batch sharding

Adds a comment on batch sharding Users sometimes encounter OOM if they run WMT on fewer machines with the same batch size, this comment explains why.