Solved by albanD in post #5 
              
              
                I guess the simplest way is to implement a custom autograd function here: 
class CheapF(autograd.Function):
  @staticmethod
  def forward(ctx, *args):
    return g(*args)

  @staticmethod
  def backward(ctx, *grad_outs):
    return f_backward(*grad_outs)

# Use it as
outputs = CheapF.apply(inputs)

N…