Matplotlib animation with different title for each frame

I have a sequence of 2D arrays which I would like to make an animation of. I would like to customize some of the frames, e.g. by adding different titles.

N_TIME_STEPS=100
records = get_records()
# records is a 250 X 250 X N_TIME_STEPS ndarray
ims = []
fig = plt.figure()
for i in range(N_TIME_STEPS):
    count_map = records[:, :, i]
    ax = plt.subplot(111)
    ax.set_title(f"step {i}")
    ax.set_ylabel(r'$k$')
    ax.set_xlabel(r'$\tau_{0}$')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    im = ax.imshow(count_map, cmap="inferno", origin='lower', extent=[0, 1, 0, 1])
    plt.colorbar(im, cax=cax)
    ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=100, blit=False, repeat_delay=1000)

This doesn't work. The title of all frames is "step 99"

Matplotlib helpfully supplies the following warning

Adding an axes using the same arguments as a previous axes currently reuses the earlier instance.  In a future version, a new instance will always be created and returned.  Meanwhile, this warning can be suppressed, and the future behavior ensured, by passing a unique label to each axes instance

But when I try this advice and use

ax = plt.subplot(111, label=str(i))

I get an unexplained

Process finished with exit code -1073740791 (0xC0000409)

What am I doing wrong?

Note: The answer to this question does not help me as blit is set to False and I want the title outside of the plot. The answer to this question is also irrelevant