EMA update for codebook

[j-min]

Is NearestEmbedEMA for EMA codebook update working?
I tried training by replacing NearestEmbed with NearestEmbed, but the loss did not converge. NearestEmbed worked for my case.

[sylqiu]

I am also curious about this. Can you try initializing cluster_size to be ones instead of zeros? Although the latter is done in the Sonnet's official implementation, I actually don't understand why it should be initialized to zeros, since in case no feature is an NN of a code, that code will seem to be multiplied by a big constant.

[nadavbh12]

Haven't tried running it with ones.
If it works for you, feel free to send a pull request.