I will take an example
class Model(nn.Module):
def __init__(self):
super(Model,self).__init__()
self.l1=nn.Linear(100,50)
self.l2=nn.Linear(50,10)
self.l3=nn.Linear(10,1)
self.sig=nn.Sigmoid()
def forward(self,x):
x=self.l1(x)
x=self.l2(x)
x=self.l3(x)
x=self.sig(x)
return(x)
class weightConstraint(object):
def __init__(self):
pass
def __call__(self,module):
if hasattr(module,'weight'):
print("Entered")
w=module.weight.data
w=w.clamp(0.5,0.7)
module.weight.data=w
# Applying the constraints to only the last layer
constraints=weightConstraint()
model=Model()
model._modules['l3'].apply(constraints)
Hope this helps