train_jax_ppo.py:ValueError: vmap got inconsistent sizes for array axes to be mapped

`python ./train_jax_ppo.py --env_name PandaPickCubeCartesian --use_tb true --visio
n true`
The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/Tu/Data/Documents/mujoco_playground/learning/./train_jax_ppo.py", line 525, in
app.run(main)
File "/home/Tu/.local/lib/python3.10/site-packages/absl/app.py", line 316, in run
_run_main(main, args)
File "/home/Tu/.local/lib/python3.10/site-packages/absl/app.py", line 261, in _run_main
sys.exit(main(argv))
File "/home/Tu/Data/Documents/mujoco_playground/learning/./train_jax_ppo.py", line 445, in main
make_inference_fn, params, _ = train_fn( # pylint: disable=no-value-for-parameter
File "/home/Tu/.local/lib/python3.10/site-packages/brax/training/agents/ppo/train.py", line 419, in train
env_state = reset_fn_(key_envs)
File "/home/Tu/.local/lib/python3.10/site-packages/mujoco_playground/_src/wrapper.py", line 129, in reset
state = self.env.reset(rng)
File "/home/Tu/.local/lib/python3.10/site-packages/brax/envs/wrappers/training.py", line 84, in reset
state = self.env.reset(rng)
File "/home/Tu/.local/lib/python3.10/site-packages/mujoco_playground/_src/wrapper.py", line 307, in reset
return self._env.reset(rng)
File "/home/Tu/.local/lib/python3.10/site-packages/mujoco_playground/_src/wrapper.py", line 174, in reset
state = jax.vmap(reset, in_axes=[self._in_axes, 0])(self._mjx_model_v, rng)
ValueError: vmap got inconsistent sizes for array axes to be mapped:
* most axes (8 of them) had size 1024, e.g. axis 0 of argument mjx_model.geom_matid of type int32[1024,84];
* one axis had size 256: axis 0 of argument rng of type uint32[256,2]