`_generate_eval_unroll` throws "INTERNAL: the requested functionality is not supported"

I am trying to run a simple SAC example for quadruped locomotion.

The script I use is
```python
env_name = 'barkour'
env = envs.get_environment(env_name)

num_timesteps = 100000

train_fn = functools.partial(
sac.train,
num_timesteps=num_timesteps,
episode_length=env.eps_length,
num_envs=1, #4096,
learning_rate=3e-4,
discounting=0.99,
batch_size=256,
num_evals=1, #10,
# normalize_observations=True,
min_replay_size=1000,
max_replay_size=num_timesteps,
network_factory=sac_networks.make_sac_networks,
randomization_fn=domain_randomize,
)

# Reset environments since internals may be overwritten by tracers from the
# domain randomization function.
env = envs.get_environment(env_name)
eval_env = envs.get_environment(env_name)

make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
```

where the environment name is the same as in the Quadruped Collab example.

However, I keep getting the error
```
E0122 13:56:59.833734 391093 pjrt_stream_executor_client.cc:3045] Execution of replica 0 failed: INTERNAL: the requested functionality is not supported
Traceback (most recent call last):
File "scripts/train_barkour_straight_sac.py", line 94, in
main()
File "scripts/train_barkour_straight_sac.py", line 67, in main
make_inference_fn, params, _ = train_fn(environment=env,
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/varun/.pyenv/versions/3.11.7/lib/python3.11/site-packages/brax/training/agents/sac/train.py", line 577, in train
metrics = evaluator.run_evaluation(
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/varun/.pyenv/versions/3.11.7/lib/python3.11/site-packages/brax/training/acting.py", line 134, in run_evaluation
eval_state = self._generate_eval_unroll(policy_params, unroll_key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the requested functionality is not supported
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

I am using `jax==0.5.0` and `brax==0.12.1`.