def _weights_init()

in torchrecipes/vision/image_generation/module/gan.py [0:0]

• 7 lines of code
• 3 McCabe index (conditional complexity)

def _weights_init(m: nn.Module) -> None:
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        # pyre-fixme[6]: Expected `Tensor` but got `Union[Tensor, nn.Module]`.
        torch.nn.init.normal_(m.weight, 0.0, 0.02)
    elif classname.find("BatchNorm") != -1:
        # pyre-fixme[6]: Expected `Tensor` but got `Union[Tensor, nn.Module]`.
        torch.nn.init.normal_(m.weight, 1.0, 0.02)
        # pyre-fixme[6]: Expected `Tensor` but got `Union[Tensor, nn.Module]`.
        torch.nn.init.zeros_(m.bias)