def train_step(model, features, targets, optimizer, loss_function):
    """Computes an optimizer step and updates NN weights.

    Parameters
    ----------
    model : NeuralNetwork instance
    features : torch.Tensor
      Input features.
    targets : torch.Tensor
      Targets to compare with in loss function.
    optimizer : torch.optim optimizer
      Optimizer to compute the step.
    loss_function : torch.nn loss function
      Loss function to compute loss.

    Returns
    -------
    loss : torch.Tensor
      The computed loss value.
    output : torch.Tensor
      The compute NN output.
    """

    model.zero_grad()
    output = model(features)
    loss = loss_function(output, targets)
    loss.backward()
    optimizer.step()

    return loss, output