Solved by albanD in post #2 
              
              
                Hi, 
You only want to save either inputs or outputs with save_for_backward. And they are saved in a slightly complex way. 
If you just want to have a state to pass around, you can give a dictionnary to your forward: 
class YourFn(Function):
  @staticmethod
    def forward(ctx, arg1, arg2, my_state):…