Thanks.
Well, I did indeed get it working to a point, but the C++ API does not seem to implement this. Firstly, retain_grad() is not defined. Secondly, when using auto cgrad = torch::autograd::grad({c}, {a, b}, {a_dispatch});, I get the following:
terminate called after throwing an instance of 'std::runtime_error'
  what():  the derivative for 'target' is not implemented

So this is also not implemented. I guess for now I will stick with what works.
Cheers.
EDIT: wait, no, this seems to be caused by requiring a gradient on b as well to satisfy autograd::grad() call, which then breaks torch::binary_cross_entropy(a, b). So that is what is not implemented. Not sure how to work around it - maybe give it another const for the time being.