Solved by ptrblck in post #2 
              
              
                The backward would expect the same number of input arguments as were returned in the forward method, so you would have to add these arguments as described in the backward section of this doc.