Denoise step in sampling

[yuhao1982]

When you do denoise from time $T$ to $T-1$ as follows:

I only see your predictor predicts $T-1$ samples stored in samples_mean, not changes batch, but you return is batch, mean_batch, recorded_samples. As you mention in previous issue, samples_mean only uses in last step. So could you explain the logic here ?

The sample or sample_mean is a ChemGraph structure. If I want to see which atomic_numbers sampled in step $T$ to $T-1$, should I just use sample["atomic_numbers"] or sample_mean["atomic_numbers"], right?