Minor Error (Probably caused by FLAX Api updates)

[stevenygd]

Fix the following error I got when running the original version of the code. `` Traceback (most recent call last): File "/home/guandao/anaconda3/envs/jax-cm/lib/python3.9/runpy.py", line 197, in _run_module_as_main return _run_code(code, main_globals, None, File "/home/guandao/anaconda3/envs/jax-cm/lib/python3.9/runpy.py", line 87, in _run_code exec(code, run_globals) File "/home/guandao/CM/consistency_models_cifar10/jcm/main.py", line 86, in app.run(main) File "/home/guandao/anaconda3/envs/jax-cm/lib/python3.9/site-packages/absl/app.py", line 308, in run _run_main(main, args) File "/home/guandao/anaconda3/envs/jax-cm/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main sys.exit(main(argv)) File "/home/guandao/CM/consistency_models_cifar10/jcm/main.py", line 66, in main train.train(FLAGS.config, FLAGS.workdir) File "/home/guandao/CM/consistency_models_cifar10/jcm/train.py", line 60, in train score_model, init_model_state, initial_params = mutils.init_model(next(rng), config) File "/home/guandao/CM/consistency_models_cifar10/jcm/models/utils.py", line 103, in init_model init_model_state, initial_params = variables.pop("params") ValueError: too many values to unpack (expected 2)

``

[LYMDLUT]

yes，I met the same problem