code: def save_state(state, output_dir, keep=False):
    
    params_file = os.path.join(output_dir, "model.pkl")
    with gfile.GFile(params_file, "wb") as f:
      pickle.dump((state.params, state.step, state.history), f)
    if keep:
      params_file = os.path.join(output_dir, "model_{}.pkl".format(state.step))
      with gfile.GFile(params_file, "wb") as f:
        pickle.dump((state.params, state.step, state.history), f)
    log("Model saved to %s" % params_file, stdout=False)
docs: Save State and optionally gin config.