Formatting issues with logging multiple images in Machine Learning Service

I have been having issues logging multiple matplotlib images in Machine Learning Service. While I am able to log the images, they frequently overlap making the image logging feature less useful. Let's look at a simple toy train.py As you can see, I use matplotlib's object oriented interface to plot 2 figures. One potential solution is to wrap all my plots into a facet grid and simply output one plot.