class_docstr: Wrapper model class to train models in distributed data parallel manner.
  The way to use this class to train your module in distributed manner is::
  
      distributed_model = DistributedModel(
          module=model,
          device_ids=[device_id0, device_id1],
          output_device=device_id0,
          broadcast_buffers=False,
      )
  
  
  where, `model` is the object of the actual model class you want to train in
  distributed manner.
class_docstr_tok: ["Wrapper","model","class","to","train","models","in","distributed","data","parallel","manner",".","The","way","to","use","this","class","to","train","your","module","in","distributed","manner","is",":",":","distributed_model","=","DistributedModel","(","module=model",",","device_ids=","[","device_id0",",","device_id1","]",",","output_device=device_id0",",","broadcast_buffers=False",",",")","where",",","`","model","`","is","the","object","of","the","actual","model","class","you","want","to","train","in","distributed","manner","."]
docstr: Override to set stage
docstr_tok: ["Override","to","set","stage"]
code: def train(self, mode=True):
      """
      Override to set stage
      """
      # use DistributedDataParallel.train since it fits distributed_training
      super().train(mode)
      self._set_module_stage(Stage.TRAIN)
  
file_path: pytext/pytext/models/distributed_model.py
imports: import torch.nn
  from pytext.common.constants import Stage
  from pytext.utils.usage import log_class_usage
label_desc: Use image node_id 6 for calling the DistributedModel obj's underlying member method code with example usage: obj.train(mode) without return types