def METHOD_NAME(self, step: int, val_step: int, epoch: int, val_batch) -> None:
    labels = val_batch["labels"]
    if len(jax.devices()) > 1:
        generated_ids = self.p_generate_step(self.state.params, val_batch)
    else:
        generated_ids = self.generate_step(self.state.params, val_batch)
    self.eval_preds.extend(jax.device_get(generated_ids.reshape(-1, 64)))
    self.eval_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))