def eval_step(batch, state, config):
  """Compute the metrics for the given model in inference mode."""

  logging.info('eval_step(batch=%s)', batch)
  variables = {'params': state.optimizer.target}
  model = models.Model(config)
  decoder_logits = model.apply(variables, batch)
  mask = model_utils.shift_right(batch['target_txt_mask'])
  mask *= batch['mask'][:, None]
  metrics = train_utils.compute_metrics(
      logits=decoder_logits,
      targets=model_utils.shift_right(batch['target_token']),
      weights=mask)
  return metrics
Compute the metrics for the given model in inference mode.
language/gscan/xattn_model/train.py