I am working on a project that requires me to write a method that is not differentiable. Hence I have calculated the gradient w.r.t the input of this method. Now I have tried to use register_hook to return the calculated grad to be able to flow the gradient backwards from there. But register_hook is not being called for this tensor. After reading through a few posts here, I realized that register_hook will not be called for tensors if gradient is not being calculated. I tried to assign gradient by doing this.
my_tensor.grad = calculated_grad. However, I do not think that has any effect.
Can you please help me on this? I am sharing a demo version of my code.
class MLP(nn.Module): def __init__(self, input_size, hidden_size, num_classes=2): super(MLP, self).__init__() self.input_size = input_size self.hidden_size = hidden_size self.fc1 = nn.Linear(self.input_size, self.hidden_size) self.relu = nn.ReLU() self.fc2 = nn.Linear(self.hidden_size,num_classes) def forward(self, x): hidden = self.fc1(x) relu = self.relu(hidden) output = self.fc2(relu) return output net = MLP(input_size, hidden_size, num_classes) out = net(data) #used out on a non-differentiable method final_output = non-diff(out) calculated_grad = grad_calc(out) # calculated gradient of out with respect to non-diff method. loss = criterion(final_output, target) loss.backward() out.grad = calculated_grad opt.step()
I would really appreciate your help on this. I have tried so many things to fix this. But nothing worked.