a minor adjustments in the colorbar

[hansen7]

Thanks for the contribution! I've find a minor adjustment is required for my dev environments.

VeloVAE/velovae/plotting.py

Line 534 in 1f918f5

cbar = plt.colorbar(sm)

should be cbar = plt.colorbar(sm, ax=ax)

[g-yichen]

Thank you for your feedback! We've added this change as suggested.