diff --git a/causal_models/hps.py b/causal_models/hps.py
index f7a63c8..a26bbce 100644
--- a/causal_models/hps.py
+++ b/causal_models/hps.py
@@ -12,12 +12,12 @@ def update(self, dict):
 embed = Hparams()
 embed.input_channels = 1
 embed.lr = 1e-3
-embed.wd = 1e-3  # 0.05
+embed.wd = 1e-3
 embed.lr_warmup_steps = 100
 embed.bottleneck = 4
 embed.cond_prior = True
 embed.z_max_res = 128
-embed.z_dim = [48, 30, 24, 18, 12, 6, 1]  # 16
+embed.z_dim = [48, 30, 24, 18, 12, 6, 1]
 embed.input_res = 224
 embed.enc_arch = "224b2d2,112b4d2,56b9d2,28b14d2,14b14d2,7b9d7,1b5"
 embed.dec_arch = "1b5,7b10,14b15,28b15,56b10,112b5,224b3"
@@ -28,16 +28,18 @@ def update(self, dict):
 embed.eval_freq = 1
 embed.grad_clip = 350
 embed.grad_skip = 25000
-embed.beta = 3.0
+embed.beta = 1.0
 embed.accu_steps = 2
 embed.bias_max_res = 64
 embed.x_like = "fixed_dgauss"
 embed.std_init = 1e-2
+embed.epochs = 20
 HPARAMS_REGISTRY["embed"] = embed
 
 padchest = copy.deepcopy(embed)
 padchest.parents_x = ["scanner", "sex"]
 padchest.context_dim = 2
+padchest.beta = 3.0
 HPARAMS_REGISTRY["padchest"] = padchest