z_h data formatting issue in EnLatentDiffusion model

[Lorenayannnnn]

Hi Minkai!

Thank you for sharing the code! Just one quick question regarding the format of h data throughout the training and sampling process.

At first, h is defined as {'categorical': one_hot, 'integer': charges} and the data is concatenated with categorical at the front of integer. However, at line 1310 of the EnLatentDiffusion model, z_h is formatted as z_h = {'categorical': torch.zeros(0).to(z_h), 'integer': z_h}, meaning that the charges part is placed before the categorical part.

Then, take sampling for instance: here z0[:, :, -1:] is used as charges, meaning z0 has a format different from that of z_h in the diffusion model.

Should z_h = {'categorical': torch.zeros(0).to(z_h), 'integer': z_h} be changed to z_h = {'categorical': z_h, 'integer': torch.zeros(0).to(z_h)} instead?

Thanks!
Tianyi