Hello,
I try to run will the custom autograd function as follows, which is used in mixed_operation
of ProxylessNAS.
class ArchGradientFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x, binary_gates, run_func, backward_func):
ctx.run_func = run_func
ctx.backward_func = backward_func
detached_x = detach_variable(x)
with torch.enable_grad():
output = run_func(detached_x)
ctx.save_for_backward(detached_x, output)
return output.data
@staticmethod
def backward(ctx, grad_output):
detached_x, output = ctx.saved_tensors
grad_x = torch.autograd.grad(output, detached_x, grad_output, only_inputs=True)
# compute gradients w.r.t. binary_gates
binary_grads = ctx.backward_func(detached_x.data, output.data, grad_output.data)
return grad_x[0], binary_grads, None, None
As shown that the gradients w.r.t. binary_gates is calculated in backward. But binary_gates.grad
is None when I track it by backward hooks. I know that asking questions for some specific project is not unwise choice. Since the authors close the issue block on Github, I have to ask you for some help, what could be the error of this case?
Thanks in advance!