[BUG] network crashes with stage=3 & dynamically-changed network

Describe the bug
In each iteration, the network is either module A + Module B or module A only. When it is stage=3, it crashes.

ds_report output

--------------------------------------------------                                                                                             
NOTE: Ops not installed will be just-in-time (JIT) compiled at                                                                                 
      runtime if needed. Op compatibility means that your system                                                                               
      meet the required dependencies to JIT install the op.   
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/opt/conda/lib/python3.8/site-packages/torch']
torch version .................... 1.9.0a0+c3d40fd
torch cuda version ............... 11.3
nvcc version ..................... 11.3
deepspeed install path ........... ['/opt/conda/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.5.8, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.9, cuda 11.3

No source files, tests, entry points, traceback, or reproduction code are provided. Start by reproducing the reported stage=3 run while alternating between module A + Module B and module A only, then use the supplied ds_report environment details to investigate the crash; done means the network completes those iterations without crashing.