config files update

mila-iqia · Jun 25, 2024 · 5ae43d6 · 5ae43d6
1 parent 61bcf4e
commit 5ae43d6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/examples/config_files/diffusion/config_diffusion_mace.yaml b/examples/config_files/diffusion/config_diffusion_mace.yaml
@@ -27,6 +27,8 @@ model:
     r_max: 5.0
     num_bessel: 8
     num_polynomial_cutoff: 5
+    num_edge_hidden_layers: 0
+    edge_hidden_irreps: 8x0e
     max_ell: 2
     interaction_cls: RealAgnosticResidualInteractionBlock
     interaction_cls_first: RealAgnosticInteractionBlock
@@ -38,10 +40,12 @@ model:
     correlation: 3
     gate: silu
     radial_MLP: [8, 8, 8]
-    radial_type: bessel
+    radial_type: gaussian
     conditional_prob: 0.0
     conditional_gamma: 2
     condition_embedding_size: 64
+    use_batchnorm: False
+    tanh_after_interaction: True
   noise:
     total_time_steps: 100
     sigma_min: 0.001  # default value

diff --git a/examples/config_files/diffusion/config_diffusion_mace_orion.yaml b/examples/config_files/diffusion/config_diffusion_mace_orion.yaml
@@ -25,8 +25,10 @@ model:
     architecture: diffusion_mace
     number_of_atoms: 8
     r_max: 5.0
-    num_bessel: 8
+    num_bessel: 'orion~choices([128, 256, 512])'
     num_polynomial_cutoff: 5
+    num_edge_hidden_layers: 0
+    edge_hidden_irreps: 8x0e
     max_ell: 2
     interaction_cls: RealAgnosticResidualInteractionBlock
     interaction_cls_first: RealAgnosticInteractionBlock
@@ -38,10 +40,12 @@ model:
     correlation: 3
     gate: silu
     radial_MLP: 'orion~choices([[8, 8, 8], [32, 32, 32], [64, 64]])'
-    radial_type: bessel
+    radial_type: 'orion~choices(["bessel", "gaussian"])'
     conditional_prob: 'orion~choices([0.0, 0.25, 0.5, 0.75])'
     conditional_gamma: 2
     condition_embedding_size: 'orion~choices([32, 64])'
+    use_batchnorm: False
+    tanh_after_interaction: True
   noise:
     total_time_steps: 100
     sigma_min: 0.001  # default value