Add tests to vector quantizer layer

[Warvito]

Fix ( #75)

[danieltudosiu]

Thank you for having a go at this. I think we are missing the following two tests:

1. Test if the EMA logic is working. This could be done by initialising the embedding weights with sequential integers filled vectors and then pick one of the integers generate a batch of images filled with that integer and apply some random gaussian noise to those image. The test would then assert if the integer filled vector from the embedding space is the only one that is changed.

2. We should also test the DDP by applying test 1 to both DDP enabled and non-DDP and asserting if the embedding dictionaries are the same.

[Warvito]

I created a issue for the distributed tests #123

[Warvito]

Test if the EMA logic is working. This could be done by initialising the embedding weights with sequential integers filled vectors and then pick one of the integers generate a batch of images filled with that integer and apply some random gaussian noise to those image. The test would then assert if the integer filled vector from the embedding space is the only one that is changed.

This is not what is implemented. I think you forgot about the decay component. In the end, all weights change. Testing the EMA logic would be a little more difficult maybe we could leave without this test for now

[danieltudosiu]

@Warvito sorry, I don't really understand what you mean. Should we take this IRL so we can create a complete test?

[Warvito]

I added the test as we discussed