Minor bug fix

kazewong · May 27, 2024 · 193ddc2 · 193ddc2
1 parent 3c52f77
commit 193ddc2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/flowMC/Sampler.py b/src/flowMC/Sampler.py
@@ -127,7 +127,7 @@ def __init__(
         self.likelihood_vec = self.local_sampler.logpdf_vmap
 
         self.optim = optax.chain(
-            optax.clip(1.0), optax.adamw(self.learning_rate, self.momentum), optax.ema(0.999)
+            optax.clip(1.0), optax.adam(self.learning_rate, self.momentum)
         )
         self.optim_state = self.optim.init(eqx.filter(self.nf_model, eqx.is_array))