output dimensions for vector quantize

[dorsa-zeinali]

Hi,
I have data that is the shape (batch, channel, height, width), I want to quantize it and I was wondering what the dimension of the vectorquantizer should be. I tried dim=height, but it produces something of dimensions [1, channel] when I need something of the dimension [batch, channel].
any help would be appreciated.
thanks

[lucidrains]

could you show a short example script? I can make the correction for you easier that way

[dorsa-zeinali]

Here is the description of my model. my code is rather long , should I include other information? although I solved the issue by setting the batch size to 1. If I train this network, will it also update the codebook with the optimizer? or will the updates just be ema ?