gradient checkpointing with "torch.utils.checkpoint()". DDP will
     work as expected when there are no unused parameters in the model
     and each layer is checkpointed at most once (make sure you are
     not passing find_unused_parameters=True to DDP). We currently
     do not support the case where a layer is checkpointed multiple
     times, or when there unused parameters in the checkpointed model.
Note:
 To let a non-DDP model load a state dict from a DDP model,
 "consume_prefix_in_state_dict_if_present()" needs to be applied
 to strip the prefix "module." in the DDP state dict before
 loading.

Warning:
 Constructor, forward method, and differentiation of the output
 (or a function of the output of this module) are distributed
 synchronization points. Take that into account in case different
 processes might be executing different code.

Warning:
 This module assumes all parameters are registered in the model by
