Set Matplotlib colorbar size to match graph

I cannot get the colorbar on imshow graphs like this one to be the same height as the graph, short of using Photoshop after the fact. How do I get the heights to match?Example of the colorbar size mismatch

You can do this easily with a matplotlib AxisDivider.

The example from the linked page also works without using subplots:

    import matplotlib.pyplot as plt
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    import numpy as np

    ax = plt.gca()
    im = ax.imshow(np.arange(100).reshape((10,10)))

    # create an axes on the right side of ax. The width of cax will be 5%
    # of ax and the padding between cax and ax will be fixed at 0.05 inch.
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)

    plt.colorbar(im, cax=cax)

enter image description here