Dimension mismatch in Residual VQ

If codebook_dim is not equal to the input embedding dimension, we get:

```
RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 128 but got size 8 for tensor number 1 in the list.

.../residual_vq.py", line 472, in forward
    all_residuals = cat((all_residuals, rearrange(residual, '... d -> ... 1 d')), dim = -2)
```