were you thinking of something specific?

Yes, I was hoping to penalize the norm of the entire Jacobian, computed in a single backwards pass. I can’t really rely on a solution that scales with the output dimension of the network. Is there no other way?
Why doesn’t autograd compute the full Jacobian ordinarily? Maybe I need to learn more about forward and backward auto-differentiation.