bug fix for layerwise params
improve memory usage efficiency
more memory optimization
Rework training and validation loss computation