adding support for kwargs in `load_state` (#989)
* adding support for kwargs in `load_state`
* Apply suggestions from code review
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
* quality ✨
* addressing comments
1. renaming variable to make it explicit
2. adding kwargs to `save_state` for parity
Co-Authored-By: Zachary Mueller <7831895+muellerzr@users.noreply.github.com>
* Apply suggestions from code review
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Zachary Mueller <7831895+muellerzr@users.noreply.github.com>
Co-authored-by: Stas Bekman <stas00@users.noreply.github.com>
