def validate_on_batch(
    network,
    loss_fn,
    X,
    y_target
):
    """Perform a forward pass on a batch of samples and compute
    the loss and the metrics.
    """
    # Do the forward pass to predict the primitive_parameters
    y_hat = network(X)
    loss = loss_fn(y_hat, y_target)
    return (
        loss.item(),
        [x.data if hasattr(x, "data") else x for x in y_hat],
    )