Passing the layer arg to the wrapper.

Fix the number of layers to be passed as an argument to the wrapper.