As far as I can tell, it’s not possible. Also see Creating a custom loss function
So, it looks like you could create a custom autograd module to handle this. If it was me, I might consider logging it on pytorch issues page and/or submitting the custom autograd module for PR.