AlphaBetaGamma96:

Would this be enable via something like this?


Yes this will work.
It won’t change the torchviz diagram. It will change the stack trace you see when it crashes.

Sorry I should’ve made this a bit clearer, the output of the custom_det function returns [B, ] for an input of [B, N, N].

So that means that the backward should expect a Tensor of size [B, ] as input and should return a Tensor of size [B, N, N]. Is that what you’re doing?

Also, I think I’ve just realised that the Tensor of [25, 4, 4] could be the grad_output which is returned in the Double Backward method?

No that would be the input that you save for backward in your custom Function most likely.