Hi
I have a loss function that looks something like this
loss = logits + lambda * get_value_of_random_fc_neuron(neuron_index)
I use forward hooks to get the random neuron value in get_value_of_random_fc_neuron()
The issue is the gradients are same if

loss = logits or
loss = logits + lambda * get_value_of_random_fc_neuron(neuron_index)

This is because the value of neuron is a tensor without grad_fn assigned. How do i make sure the gradient is calculated for full loss.
Thanks in advance