dimension issue on training

[dgm2]

Hi great work, really nice.
I am getting an error in training, not sure if you have seen this before or know a quick fix / possible cause ?

this line
https://github.com/gracezhao1997/EEGSDE/blob/master/energys_prediction/en_diffusion.py#L380
not sure why xh has a dimension 7 instead of 8 as eps.

z_t = alpha_t * xh + sigma_t * eps
RuntimeError: The size of tensor a (7) must match the size of tensor b (8) at non-singleton dimension 2

running the following

run_train_property_prediction_energy.py
--exp_name=predict_mu
--conditioning=mu
--model=egnn_dynamics
--lr=1e-4
--nf=192
--n_layers=7
--weight_decay=1e-16
--save_model=True
--diffusion_steps=3
--sin_embedding=False
--n_epochs=1
--diffusion_noise_schedule=polynomial_2
--diffusion_noise_precision=1e-5
--dequantization=deterministic
--include_charges=False
--normalize_factors=[1,8,1]
--dataset=qm9_second_half

best regards

[dgm2]

it bad preprocessing from my side I think

[gracezhao1997]

Hi, the xh should have a dimension 8, where x is the molecule position with dimension 3 and h is the one hot representation of atom types with dimension 5 (C,H,O,F,N). I recheck the provided code and it can correctly train the property prediction network used in energy function. I think you should check the dataloader.