def training_loss_fn(params):
    variables = {'params': params, **train_state.model_state}
    (output, new_model_state) = flax_model.apply(variables, batch, mutable=['batch_stats'], train=True, rngs={'dropout': dropout_rng}, debug=debug)
    loss = loss_fn(output, batch, variables['params'])
    return (loss, (new_model_state, output))
def training_loss_fn(params):
    variables = {'params': params, **train_state.model_state}
    (output, new_model_state) = flax_model.apply(variables, batch, mutable=['batch_stats'], train=True, rngs={'dropout': dropout_rng}, debug=debug)
    loss = loss_fn(output, batch, variables['params'])
    return (loss, (new_model_state, output))