Function zero_module

monai/networks/nets/diffusion_model_unet.py:51–57

Zero out the parameters of a module and return it.

(module: nn.Module)

def zero_module(module: nn.Module) -> nn.Module:
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


class DiffusionUNetTransformerBlock(nn.Module):