Within the forward and backward of an autograd.Function, autograd tracing is disabled by default (similar to when you do with torch.no_grad():), so aux_loss does not require gradient. If you wrap the aux_loss with with torch.enable_grad(): your code seems to run. (You don’t return anything from the backward, though, and it you might have funny side effects if you use backward that propagates to portions of the graph outside the function’s variables in your Function.)
Best regards
Thomas