def init_module(module, init_weight=None, init_bias=None):
    """
    Initialize the given module using the given initialization scheme for the weight and bias terms.

    The user can select the initialization scheme from `torch.nn.init`. This function is taken from [1] and modified.

    Args:
        module (torch.nn.Module): torch module to initialize.
        init_weight (callable, None): this is a callable function that accepts as input the module to initialize its
            weights. If None, it will leave the module weights untouched.
        init_bias (callable, None): this is a callable function that accepts as input the module to initialize its
            bias terms. If None, it will leave the module bias untouched.

    Returns:
        torch.nn.Module: initialized module.

    References:
        [1] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/utils.py
    """
    if init_weight is not None:
        init_weight(module.weight.data)
    if init_bias is not None:
        init_bias(module.bias.data)
    return module