Fix a bug that the lambda for second layer was not used