Given the stack trace, it seems that the error comes from a mul op that has Tensors with different sizes.
So you’re most likely doing a computation with things that have the wrong size in your custom backward or returning something that does not have the expected size (but this one should rise a nice error).
Could you get a small code sample of 30/40 lines that reproduces the crash?