Solved by tom in post #2 
              
              
                Similar to your other question, you need with torch.enable_grad(): if you want to record the graph within the forward or backward of a autograd.Function. 
Best regards 
Thomas