Solved by tom in post #2 
              
              
                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 …