import torch


def _train_step_mlp(x_batch, y_batch, model, optimizer):
    """
    Perform single standard training step.

    Args:
        x_batch:    Batch of inputs
        y_batch:    Batch of labels
        model:      torch.nn.Module
        optimizer:  torch.optim.Optimizer for model parameters

    Returns:
        Model output, batch loss
    """
    # Compute forward pass on current batch
    output = model(x_batch)

    # Compute batch loss as cross entropy loss
    batch_loss = torch.nn.functional.cross_entropy(output, y_batch)

    # Compute gradients wrt current batch loss and perform parameter update
    optimizer.zero_grad()
    batch_loss.backward()
    optimizer.step()

    return output, batch_loss