Added functionality to the save feature