Spaces:

flax-community
/

dalle-mini

Running

boris commited on Jul 16, 2021

Commit

f65ccb3

1 Parent(s): a173dad

fix: missing arg

Former-commit-id: ae6b730efc6819a366c2b35810d0f2c0a90b3227

Files changed (1) hide show

seq2seq/run_seq2seq_flax.py CHANGED Viewed

@@ -703,7 +703,7 @@ def main():
     )
     if model_args.from_checkpoint is not None:
         # restore optimizer state, step and optimizer_step
-        restore_state(state)
     # label smoothed cross entropy
     def loss_fn(logits, labels):

     )
     if model_args.from_checkpoint is not None:
         # restore optimizer state, step and optimizer_step
+        restore_state(state, artifact_dir)
     # label smoothed cross entropy
     def loss_fn(logits, labels):