Save optimizer state and dataset state
Save optimizer state and dataset state