def METHOD_NAME(self, params, state, batch, dropout_rng, labels):
    logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
    loss = self.loss_helper(logits, labels, batch)
    return loss