Hi,
Thanks for clearing that up.
I have another question if you could please help me–
i am trying to apply a kmeans quantization technique to the weights of each layer during the training procedure.
What i mean is during the forward propagation at each layer i want to first use the kmeans algorithm to calculate the weights and then use these calculated weights and discard the old ones.
Similarly the same procedure for the backpropagation step also.
I am trying to use hooks to implement this functionality and for some trial purposes just to get an idea of how a hook works i have written the below code-
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.cl1 = nn.Linear(5, 4)
        self.cl2 = nn.Linear(4, 2)
        
    def forward(self, x):
        print(model.cl1.weight,model.cl1.bias)
        x = self.cl1(x)
        print(x)
        x = F.relu(self.cl2(x))
        return x


activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
        print(model.weight)
        model.weight.data=torch.ones(4,5)
        print(model.weight)
        output=torch.mm(input[0],model.weight.T)
        return output
    return hook

model = MyModel()
model.cl1.register_forward_hook(get_activation('fc2'))
x = torch.ones(1, 5)
output = model(x)

and the output i am getting is as follows–
Parameter containing:
tensor([[-0.3525,  0.4128,  0.3498,  0.0078, -0.3536],
        [-0.2505, -0.0612, -0.1943, -0.2778, -0.3931],
        [-0.2942,  0.0659,  0.0260,  0.0492, -0.2993],
        [ 0.1888,  0.2567,  0.0124,  0.3343, -0.3156]], requires_grad=True) Parameter containing:
tensor([-0.3972, -0.0496, -0.0852, -0.1048], requires_grad=True)
Parameter containing:
tensor([[-0.3525,  0.4128,  0.3498,  0.0078, -0.3536],
        [-0.2505, -0.0612, -0.1943, -0.2778, -0.3931],
        [-0.2942,  0.0659,  0.0260,  0.0492, -0.2993],
        [ 0.1888,  0.2567,  0.0124,  0.3343, -0.3156]], requires_grad=True)
Parameter containing:
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]], requires_grad=True)
tensor([[5., 5., 5., 5.]], grad_fn=<MmBackward>)

As you can see from the above output i am getting the value but the grad_fn of the variable x changes.
What is the best way for me to overwrite the weights of each layer while not affecting the computation graph?
Thanks!