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.