Update plot.py

Figures should be explicitly closed after saving for 1D and 2D plots in order to avoid memory issues/errors.