Hi all, I have a computational graph x -> y -> loss.
I want to compute d_loss/d_y, but NOT d_loss/d_x
In theano I’d use disconnected_grad 54.
How can I do this in PyTorch?
I’ve heard you can somehow use the register_backward_hook method, but I have not found an example of this working.  I’ve attached an attempt below, but it seems that the backward hook function, which is supposed to block the gradient computation, is never called.
import torch
import numpy as np

"""
We have the following equations:

x = 3
y = 4*x
loss = (v2-10)**2

This should yield
d_loss_d_y == 2*(3*4-10) == 4
d_loss_d_x == d_loss_d_y*4 == 16

Now - we ONLY want to calcuate d_loss_d_y.  We want to block the gradient backpropagation
so that we don't bother calculating d_loss_d_x.

"""

x = torch.autograd.Variable(torch.from_numpy(np.array([3.])), requires_grad=True)

class MyOp(torch.nn.Module):

    def __call__(self, x):
        return x*4

op = MyOp()

def my_hook(mod, in_grad, out_grad):
    print 'Backward Hook Called (THIS NEVER HAPPENS)'
    return None  # (Should stop further backpropagation?)

op.register_backward_hook(my_hook)
y = op(x)

intermediate_grads = {}

y.register_hook(lambda grad: intermediate_grads.setdefault(y, grad))

loss = ((y - 10.) ** 2).sum()
loss.backward()


assert np.array_equal(intermediate_grads[y].data.numpy(), [4.])

# The following should raise some kind of exception .. because d_loss_d_x shouldn't be computed
assert not np.array_equal(x.grad.data.numpy(), [16.])