diff --git a/config/experiments/ccube/corners.yaml b/config/experiments/ccube/corners.yaml
index e3594ac76..ccc207c6f 100644
--- a/config/experiments/ccube/corners.yaml
+++ b/config/experiments/ccube/corners.yaml
@@ -40,18 +40,17 @@ gflownet:
     z_dim: 16
     lr_z_mult: 100
     n_train_steps: 10000
-  policy:
-    forward:
-      type: mlp
-      n_hid: 512
-      n_layers: 5
-      checkpoint: forward
-    backward:
-      type: mlp
-      n_hid: 512
-      n_layers: 5
-      shared_weights: False
-      checkpoint: backward
+
+# Policy
+policy:
+  forward:
+    type: mlp
+    n_hid: 128
+    n_layers: 2
+    checkpoint: forward
+  backward:
+    shared_weights: True
+    checkpoint: backward
 
 # WandB
 logger:
diff --git a/config/experiments/ccube/uniform.yaml b/config/experiments/ccube/uniform.yaml
index 6970a3e95..a81d58d05 100644
--- a/config/experiments/ccube/uniform.yaml
+++ b/config/experiments/ccube/uniform.yaml
@@ -40,18 +40,17 @@ gflownet:
     z_dim: 16
     lr_z_mult: 100
     n_train_steps: 10000
-  policy:
-    forward:
-      type: mlp
-      n_hid: 256
-      n_layers: 3
-      checkpoint: forward
-    backward:
-      type: mlp
-      n_hid: 256
-      n_layers: 3
-      shared_weights: False
-      checkpoint: backward
+
+# Policy
+policy:
+  forward:
+    type: mlp
+    n_hid: 128
+    n_layers: 2
+    checkpoint: forward
+  backward:
+    shared_weights: True
+    checkpoint: backward
 
 # WandB
 logger: