Adapt the following method. This will also require further adaption of the corruption module, that's why it was ignored so far. The current version of the TD predictor-corrector is creating a lot of warnings, as the noise batches are indeed already expanded (since the timestep is a site property rather than a structural property as in the mattergen model). To avoid confusion, it would be good and also the cleaner solution to actually adapt it accordingly.
https://github.com/microsoft/mattergen/blob/a245cf2b7538eea6d873e6430b0e30c56d26c60e/mattergen/diffusion/corruption/corruption.py#L40-L55
Adapt the following method. This will also require further adaption of the corruption module, that's why it was ignored so far. The current version of the TD predictor-corrector is creating a lot of warnings, as the noise batches are indeed already expanded (since the timestep is a site property rather than a structural property as in the mattergen model). To avoid confusion, it would be good and also the cleaner solution to actually adapt it accordingly.
https://github.com/microsoft/mattergen/blob/a245cf2b7538eea6d873e6430b0e30c56d26c60e/mattergen/diffusion/corruption/corruption.py#L40-L55