def train_step(model, optimizer, loss_fn, conditions, true, out):
    """One training step
    
    Args:
        model: the feed-forward network
        optimizer: the optimizer for the network
        loss_fn: the loss function
        conditions: the observing conditions used as inputs
        true: the true galaxy magnitudes used as inputs
        out: the ground truth output
    """

    optimizer.zero_grad()
    conditions.requires_grad_(True)
    true.requires_grad_(True)
    
    predout = model(conditions, true).squeeze()
    loss = loss_fn(predout, out)
    loss.backward()
    optimizer.step()
    return loss.item(), predout.data