File size: 264 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 | from torch import nn
from einops import rearrange, parse_shape
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
|