Solved by albanD in post #2 
              
              
                No this is expected. You should return as many elements as there were inputs to the forward function. If something is non-differentiable, you should return None.