class_docstr: This class mirrors deepspeed functionality. Not all changes are implemented in this version.
  
  See TorchEngineFull for more modifications.
class_docstr_tok: ["This","class","mirrors","deepspeed","functionality",".","Not","all","changes","are","implemented","in","this","version",".","See","TorchEngineFull","for","more","modifications","."]
code: def _init_distributed(self, model):
      model = torch.nn.parallel.DistributedDataParallel(
          model,
          device_ids=[self.setup["device"]]
          if self.setup["device"].type == "cuda"
          else None,
          output_device=self.setup["device"]
          if self.setup["device"].type == "cuda"
          else None,
          broadcast_buffers=self.cfg_impl.broadcast_buffers,
          bucket_cap_mb=self.cfg_impl.bucket_cap_mb,
          gradient_as_bucket_view=self.cfg_impl.gradient_as_bucket_view,
          static_graph=self.cfg_impl.static_graph,
      )
      return model
  
imports: import torch
  import torch._inductor.utils
  import os
  import json
  from omegaconf import OmegaConf
  from functools import partial
  from contextlib import nullcontext
  import time
  import logging
  import transformers
  from safetensors.torch import load_file, save_file
  from transformers.utils.generic import working_or_temp_dir
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from .utils import group_parameters, prepare_pretraining_dataloader, update_ema, updated_latest_weight_average
  from .optimizers.schedulers import get_schedule_fn
  from .optimizers import Adahessian, AdamWScale, Shampoo, LARS, SAM, ProgressiveBatching, AGD, Sophia
  import warnings
label_desc: Use image node_id 12 for calling the TorchEngineMinimal obj's underlying member method code with example usage: obj._init_distributed(model) and returns: model