Some of the relevant snippets of code:
    video_array = np.zeros(
        shape=(1, 3, trajectory_len, height*100, width*100),
        dtype=np.uint8)

    for trajectory_idx in trajectory_indices:

        fig, axes = plt.subplots(
            1,
            2,
            figsize=(width, height),
            gridspec_kw={'width_ratios': [1, 0.05]})
        fig.suptitle('Example Trajectory')
        # plot the first trajectory
        sc = axes[0].scatter(
            x=[x[trajectory_idx]],
            y=[y[trajectory_idx]],
            c=[trajectory_indices[trajectory_idx]],
            s=4,
            vmin=0,
            vmax=trajectory_len,
            cmap=plt.cm.jet)

        plt.close(fig=fig)

        video_array[0, :, trajectory_idx, :, :] = np.transpose(data, (2, 0, 1))

    hook_input['tensorboard_writer'].add_video(
        tag='sampled_trajectory',
        vid_tensor=torch.from_numpy(video_array),
        global_step=hook_input['grad_step'],
        fps=4)


def fig2data(fig):

    # draw the renderer
    fig.canvas.draw()

    # Get the RGB buffer from the figure
    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    return data