def setup_horovod(model, learning_rate):
    """ Setup for Horovod usage.

    Args:
        model(MultitaskModel): The MultitaskModel object.
        learning_rate(float): Learning rate for the model.

    Returns: hvd.DistributedOptimizer: Optimizer to use for computing
    gradients and applying updates.

    """
    # Horovod: scale learning rate by the number of GPUs.
    optimizer = optim.Adam(model.parameters(), lr=learning_rate * hvd.size())

    # Horovod: broadcast parameters & optimizer state.
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    hvd.broadcast_optimizer_state(optimizer, root_rank=0)

    # Horovod: (optional) compression algorithm.
    compression = hvd.Compression.fp16

    # Horovod: wrap optimizer with DistributedOptimizer.
    optimizer = hvd.DistributedOptimizer(
        optimizer,
        named_parameters=model.named_parameters(),
        compression=compression,
    )

    return optimizer
 Setup for Horovod usage.

    Args:
        model(MultitaskModel): The MultitaskModel object.
        learning_rate(float): Learning rate for the model.

    Returns: hvd.DistributedOptimizer: Optimizer to use for computing
    gradients and applying updates.

    
examples/sentence_similarity/gensen_train.py