Autograd won’t be able to create the computation graph for the numpy opertations, so you would have to write a custom autograd.Function as described here 15 and implement the backward method manually.