Dreambooth: save / restore training state (#1668)
* Dreambooth: save / restore training state.
* make style
* Rename vars for clarity.
Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
* Remove unused import
Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
