Added a parameter to the pretraining script to allow specific layer unit configurations on the command line.