vae.py line 320 use this to plot

[djaym7]

replace these lines - plt.scatter(m1s, m2s, c=labels) plt.legend([str(i) for i in set(labels)])

with the lines below - scatter = plt.scatter(m1s, m2s, c=labels) plt.legend(handles=scatter.legend_elements()[0], labels=set(labels))