examples/research_projects/jax-projects/big_bird/bigbird_flax.py/Trainer/evaluate
class Trainer:    def evaluate(self, state, dataset):
        dataloader = get_batched_dataset(dataset, self.args.batch_size)
        total = len(dataset) // self.args.batch_size
        running_loss = jnp.array(0, dtype=jnp.float32)
        i = 0
        for batch in tqdm(dataloader, total=total, desc="Evaluating ... "):
            batch = self.data_collator(batch)
            metrics = self.val_step_fn(state, **batch)
            running_loss += jax_utils.unreplicate(metrics["loss"])
            i += 1
        return running_loss / i