import torch


def train(data, model, loss_fun, num_epochs=100, verbose=True, lr=0.01, logger=lambda loss: None):
    """
    train an embedding model

    Args:
        data: network data
        model: embedding auto-encoder model
        loss_fun: loss function to use with model (takes arguments ``model``, ``data``)
        num_epochs: number of training epochs
        verbose: if ``True``, display training progress (default: ``True``)
        lr: learining rate (default: 0.01)
        logger: function that receives the training loss as input and is called after each epoch (does nothing by default)

    Returns:
        trained model

    This function uses the Adam optimizer for training.
    """
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for e in range(num_epochs):
        optimizer.zero_grad()
        loss = loss_fun(model, data)
        loss.backward()
        optimizer.step()
        logger(float(loss))
        if verbose:
            print(f'epoch {e}: loss={loss.item()}')
        # schedule.step()
    return model