Improvements to the model training Lib in Python

Over time we will add several improvements to the current python library that is present in the Sagemaker Training Docker images, and is used for dataser process pipeline and custom model architecture.