weight_scale

class torchelie.nn.utils.weight_scale(module: torch.nn.modules.module.Module, name: str = 'weight', scale: float = 0)

Multiply getattr(module, name) by scale on forward pass as a hook. Used to implement equalized LR for StyleGAN