Addressing Issue #5241: Updating deepspeed/runtime/zero/stage_1_and_2.py

Addressing Issue #5241: Updating deepspeed/runtime/zero/stage_1_and_2.py #12496