mild refactor of set_requires_grad
@@ -65,6 +65,10 @@ class IdentityModule(nn.Module):
def set_requires_grad(named_parameters, requires_grad):
for name, param in named_parameters:
+ set_requires_grad_single(param, requires_grad)
+
+
+def set_requires_grad_single(param, requires_grad):
param.requires_grad = requires_grad
