Merge pull request #105 from Qianruipku:hotfix2

PiperOrigin-RevId: 807836181
Change-Id: Id83ca6c48188051d92c88d3983ae8712e2186864

Author: jsspencer

ferminet/train.py

@@ -883,7 +883,8 @@ def learning_rate_schedule(t_: jnp.ndarray) -> jnp.ndarray:
   elif isinstance(optimizer, optax.GradientTransformation):
     # optax/optax-compatible optimizer (ADAM, LAMB, ...)
     opt_state = jax.pmap(optimizer.init)(params)
-    opt_state = opt_state_ckpt or opt_state  # avoid overwriting ckpted state
+    if opt_state_ckpt is not None:
+      opt_state = tuple(opt_state_ckpt)
     step = make_training_step(
         mcmc_step=mcmc_step,
         optimizer_step=make_opt_update_step(evaluate_loss, optimizer),