DeepSpeed

db98cc3a - Fix assertion for offloading states (#6855)

Fix assertion for offloading states (#6855) This PR fixes the assertions in `offload_states` method mentioned in #6833. Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>