Colorbars not working when using make_axis_locatable

[theRealSuperMario]

trafficstars

After checking the related issues to colorbars (#111 , #115 ), I found that the colorbar is not working when using this snippet https://stackoverflow.com/questions/32462881/add-colorbar-to-existing-axis

Is there a quick fix for it?

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

RC_PARAMS = {"figure.figsize": [5, 5], "figure.dpi": 220, "pgf.rcfonts": False}

data = np.zeros((3, 3))
data[:2, :2] = 1.0

with plt.rc_context(rc=RC_PARAMS):
    fig, ax = plt.subplots(1, 1, figsize=(5, 5))
    im = ax.imshow(data)
    ax.set_title("Works")
    ax.set_axis_off()
    plt.colorbar(im, ax=ax, orientation="horizontal")
    matplotlib2tikz.save("test1.tikz")

plt.close()

with plt.rc_context(rc=RC_PARAMS):
    fig, ax = plt.subplots(1, 1, figsize=(5, 5))
    im = ax.imshow(data)
    ax.set_title("Doesn't")
    ax.set_axis_off()
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    plt.colorbar(im, cax=cax)
    matplotlib2tikz.save("test2.tikz")