class_docstr: A :class:`allennlp.nn.parallel.ddp_accelerator.DdpAccelerator` for FairScale's [`FullyShardedDataParallel`]
  (https://fairscale.readthedocs.io/en/latest/api/nn/fsdp.html).
  
  To save memory while initializing a model, you should call [`.wrap_module()`](#wrap_module) on submodules
  as they're created.
  
  See the :class:`allennlp.modules.transformer.t5.T5` class for an example of how to use this.
class_docstr_tok: ["A",":","class",":","`","allennlp.nn.parallel.ddp_accelerator.DdpAccelerator","`","for","FairScale","'s","[","`","FullyShardedDataParallel","`","]","(","https",":","\/\/fairscale.readthedocs.io\/en\/latest\/api\/nn\/fsdp.html",")",".","To","save","memory","while","initializing","a","model",",","you","should","call","[","`",".wrap_module","(",")","`","]","(","#","wrap_module",")","on","submodules","as","they","'re","created",".","See","the",":","class",":","`","allennlp.modules.transformer.t5.T5","`","class","for","an","example","of","how","to","use","this","."]
code: def wrap_model(
      self, model: "Model"
  ) -> Tuple["Model", DdpWrappedModel]:
      wrapped_model = _FSDP(
          model,
          **self._fsdp_kwargs,
      )
      if not self._fsdp_kwargs[
          "mixed_precision"
      ] and self.cuda_device != torch.device("cpu"):
          wrapped_model = wrapped_model.cuda()
      # `_FSDP._lazy_init()` may have been called already on submodules that were wrapped
      # (through `wrap_module()`), leading those submodules to think they are root submodules.
      # So we need to call `_FSDP._reset_lazy_init()` on any of these now.
      for module in wrapped_model.modules():
          if isinstance(module, _FSDP):
              module._reset_lazy_init()
      return model, FairScaleFsdpWrappedModel(
          wrapped_model,
          local_rank=self.local_rank,
          world_size=self.world_size,
      )
  
imports: import os
  from typing import Tuple, Union, Optional, TYPE_CHECKING, List, Any, Dict, Sequence
  from fairscale.nn import FullyShardedDataParallel
  from fairscale.nn.wrap import enable_wrap, wrap
  from fairscale.nn.misc import FlattenParamsWrapper
  from fairscale.optim.grad_scaler import GradScaler
  import torch
  from torch.cuda import amp
  from allennlp.nn.parallel.sharded_module_mixin import ShardedModuleMixin
  from allennlp.nn.parallel.ddp_accelerator import DdpAccelerator, DdpWrappedModel, StateDictType, LoadStateDictReturnType
label_desc: Use image node_id 2 for calling the FairScaleFsdpAccelerator obj's underlying member method code with example usage: obj.wrap_model(model) and returns: model,  FairScaleFsdpWrappedModel