def scale_module(module, scale):
    """
    Scale the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().mul_(scale)
    return module

    Scale the parameters of a module and return it.
    
trellis/modules/utils.py