[BUGFIX] Fix the save function