SageMaker fix
This removes some code that was accidentally checked-in which made some modules broken; it also optimizes the imports of SageMaker files.
This removes some code that was accidentally checked-in which made some modules broken; it also optimizes the imports of SageMaker files.