Add missing context placement for QM9 train-time sampling

Author: amorehead

.gitignore

@@ -173,8 +173,8 @@ outputs/
 epoch_*/
 .cache/
 
-data/EDM/GEOM
-data/EDM/QM9
+data/EDM/GEOM*
+data/EDM/QM9*
 
 # NFS
 .nfs*

src/models/qm9_mol_gen_ddpm.py

@@ -783,7 +783,11 @@ def sample_and_analyze(
             assert int(num_nodes.max()) <= max_num_nodes
 
             # context-conditioning
-            context = None
+            if self.condition_on_context:
+                if context is None:
+                    context = self.props_distr.sample_batch(num_nodes)
+            else:
+                context = None
 
             xh, batch_index, _ = self.ddpm.mol_gen_sample(
                 num_samples=num_samples_batch,