fig, ax = plt.subplots(figsize=(5, 5))

def plot_generated_images_draw(...):

    plt.clf()
    for i in range(generated_images.shape[0]):
        plt.subplot(dim[0], dim[1], i + 1)
        plt.imshow(generated_images[i], cmap='gray_r')
        plt.axis('off')
    if epoch == 1 or epoch % 5 == 0:
        plt.savefig('gan_generated_image_epoch_%d.png' % epoch)
    plt.pause(0.1)