DDP: RuntimeError: shape mismatch: value tensor of shape [535, 256] cannot be broadcast to indexing result of shape [537, 256]

 hi @lucidrains 
I found the same error #142 
this is my code
`self.vq = VectorQuantize(
            dim=in_channels,
            codebook_size=codebook_size,
            codebook_dim=codebook_dim,
            threshold_ema_dead_code=threshold_ema_dead_code,
            kmeans_init=False,
        )`