Activation gradient penalty

Hi @michaelklachko, if you were able to resolve this issue, can you take a look at Optimizing a loss based on autograd.grad output? I think I have a similar problem.