kinda works now

dirmeier · Feb 28, 2024 · d7e16ef · d7e16ef
1 parent 41dd5b0
commit d7e16ef
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sbijax/_src/scmpe.py b/sbijax/_src/scmpe.py
@@ -176,7 +176,7 @@ def _fit_model_single_round(
 
  @jax.jit
  def ema_update(params, avg_params):
- return optax.incremental_update(params, avg_params, step_size=0.001)
+ return optax.incremental_update(avg_params, params , step_size=0.01)
 
  @jax.jit
  def step(params, ema_params, rng, state, **batch):
@@ -189,7 +189,7 @@ def step(params, ema_params, rng, state, **batch):
  return loss, new_params, new_ema_params, new_state
 
  losses = np.zeros([n_iter, 2])
- early_stop = EarlyStopping(1e-3, n_early_stopping_patience)
+ early_stop = EarlyStopping(1e-3, n_early_stopping_patience*2)
  best_params, best_loss = None, np.inf
  logging.info("training model")
  for i in tqdm(range(n_iter)):