InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (1, 596, 256) and (1, 24, 24, -1)

[Al-aminI]

array = a python list with integers of length 596,

note: the array was not encoded by the vqgan, i just make them up, and converted them to jax jnp, but figured that the indices

returned by the vqgan encode has some additional properties

jax_arr = jnp.array(array], dtype=jnp.int32)

reconstruct from indices indices

reconstructed_img =vq_model.decode_code(jax_arr)

here is the full trace:
`## or to reconstruct using indices
---->
File [c:\Users](file:///C:\codes\vqgan-jax\vqgan_jax\modeling_flax_vqgan.py:643, in VQGANPreTrainedModel.decode_code(self, indices, params)
642 def decode_code(self, indices, params: dict = None):
--> 643 return self.module.apply({"params": params or self.params},
644 jnp.array(indices, dtype="i4"),
645 method=self.module.decode_code)

[... skipping hidden 6 frame]

File [c:\Users](file:///C:/Users/codes\vqgan-jax\vqgan_jax\modeling_flax_vqgan.py:566, in VQModule.decode_code(self, code_b)
565 def decode_code(self, code_b):
--> 566 hidden_states = self.quantize.get_codebook_entry(code_b)
567 hidden_states = self.decode(hidden_states)
568 return hidden_states

[... skipping hidden 2 frame]

File [c:\Users](file:///C:/Users/\codes\vqgan-jax\vqgan_jax\modeling_flax_vqgan.py:526, in VectorQuantizer.get_codebook_entry(self, indices, shape)
524 batch, num_tokens = indices.shape
525 z_q = self.embedding(indices)
...
1901 if isinstance(r, Tracer) or r != 0:
-> 1902 raise InconclusiveDimensionOperation(f"Cannot divide evenly the sizes of shapes {tuple(s1)} and {tuple(s2)}")
1903 return q

InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (1, 596, 256) and (1, 24, 24, -1)`