import torch


def _train_step_ddc(x_batch, y_batch, model):
    """
    Single training step for DDC.

    Args:
        x_batch:    Batch of inputs
        y_batch:    Batch of labels
        model:      ddc.DynamicDropConnect object

    Returns:
        Model output, batch loss
    """

    # Compute forward pass on current batch
    output = model(x_batch)

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

    # Compute gradients of model parameters wrt current batch and detach
    gradients = torch.autograd.grad(batch_loss, model.parameters())
    gradients = [g.detach() for g in gradients]

    # Optimize all parameters
    model.update_params(gradients)

    # Apply weight probability update
    model.update_probs(gradients)

    return output, batch_loss