Tell Pytorch's Autograd not to include a operation for gradient computation

[ptrblck]

Both approaches would break the computation graph and the previous layers wouldn’t get a valid gradient.
You could try to adapt this workflow (or write a custom autograd.Function and define the backward method manually).