Set requies_grad for all the networks. Args: nets (nn.Module | list[nn.Module]): A list of networks or a single network. requires_grad (bool): Whether the networks require gradients or not
(nets, requires_grad=False)
| 33 | |
| 34 | |
| 35 | def set_requires_grad(nets, requires_grad=False): |
| 36 | """Set requies_grad for all the networks. |
| 37 | |
| 38 | Args: |
| 39 | nets (nn.Module | list[nn.Module]): A list of networks or a single |
| 40 | network. |
| 41 | requires_grad (bool): Whether the networks require gradients or not |
| 42 | """ |
| 43 | if not isinstance(nets, list): |
| 44 | nets = [nets] |
| 45 | for net in nets: |
| 46 | if net is not None: |
| 47 | for param in net.parameters(): |
| 48 | param.requires_grad = requires_grad |
| 49 | |
| 50 | |
| 51 | def zero_module(module): |