Jupyter title bug in for loop

Question

I am making a scatter plot with python that updates with a for loop and I am having a issue were my title moves down to the bottom of the figure. I've used this code in PyCharm and there it works just fine, but when I move it to Jupyter it breaks.

I did add %matplotlib qt to make sure that an external figure appears.

This is a part of the code I am using in Jupyter

%matplotlib qt

def md_sim():
    pos_init()
    pos_random()

    x = []
    y = []
    z = []
    for i in range(0, M):
        x.append(xyz[i][0])
        y.append(xyz[i][1])
        z.append(xyz[i][2])
    plt.ion()
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    data = ax.scatter(x, y, z, s=100)
    ax.set_title("Initial lattice positions")
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    fig.show()

    vel_random_gauss()
    correct_v()
    LJ_pot(0)
    for i in range(0, steps):
        verlet(i+1)
        correct_v()

        x = []
        y = []
        z = []
        for n in range(0, M):
            x.append(xyz[n][0])
            y.append(xyz[n][1])
            z.append(xyz[n][2])
        plt.pause(0.001)
        data._offsets3d = (x, y, z)
        ax.set_title("Lattice positions after " + str(i) + " steps.")
        plt.draw()

I have made a small video of the title moving down in the figure https://gyazo.com/e203cf86974ee53edeb4534edf443a03

Please help me out!