def train_network(model, data, epochs, loss_function, optimizer, display_loss=True):
    """
    Train a pytorch model.
    Inputs:
       model: torch model, e.g., MLPRegressor()
       data: dict, contains at a minimum:
          {"x_train": <torch.Tensor>,
           "y_train": <torch.Tensor>,
           "x_val": <torch.Tensor>,
           "y_val": <torch.Tensor>}
        epochs: int, number of training epochs
        loss_function: torch loss function, e.g., torch.nn.MSELoss()
        optimizer: torch optimizer, e.g., torch.optim.Adam()
        display_loss: bool; If True, display the training/validation loss at
           each epoch. Defaults to True.
    Outputs:
       model: trained torch model
    """

    for _ in range(epochs):
        y_pred = model.forward(data["x_train"])
        loss = loss_function(y_pred, data["y_train"])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # compute validation loss
        y_pred_val = model(data["x_val"])
        val_loss = loss_function(y_pred_val, data["y_val"])
        if display_loss:
            print("training/validation loss",
                  round(loss.item(), 3),
                  "/",
                  round(val_loss.item(), 3))

    return model