def init_keras(hvd=None):
    """
    Set config for Horovod. Config params copied from official example:
    https://github.com/uber/horovod/blob/master/examples/keras_mnist_advanced.py#L15

    :param hvd: instance of horovod.keras
    """

    init_cuda_env()
    config = tf.ConfigProto()

    if hvd:
        hvd.init()
        config.gpu_options.allow_growth = True
        config.gpu_options.visible_device_list = str(hvd.local_rank())

    set_session(tf.Session(config=config))

    Set config for Horovod. Config params copied from official example:
    https://github.com/uber/horovod/blob/master/examples/keras_mnist_advanced.py#L15

    :param hvd: instance of horovod.keras
    
cakechat/utils/env.py