class_docstr: Manages a PyTorch model for training.
class_docstr_tok: ["Manages","a","PyTorch","model","for","training","."]
code: def setup_distributed(
      self, mode, cluster_info, driver_ip, driver_tcp_store_port
  ):
      if mode == "fit":
          self.rank = get_rank(cluster_info)
          logger.info(f"cluster is: {cluster_info}")
          self.setup_components()
          self.setup_torch_distribute(
              tcp_store_host=driver_ip,
              tcp_store_port=driver_tcp_store_port,
              world_rank=self.rank,
              world_size=self.size,
          )
      else:
          self.rank = 0
          self.setup_components()
          if self.model_creator:
              self.setup_operator(self.models)
  
file_path: BigDL/python/orca/src/bigdl/orca/learn/pytorch/pytorch_pyspark_worker.py
imports: from bigdl.orca.learn.pytorch.torch_runner import TorchRunner
  import torch.distributed
  import logging
  import os
  import tempfile
  import copy
  from pyspark import BarrierTaskContext, TaskContext
  from bigdl.orca.learn.utils import save_pkl, duplicate_stdout_stderr_to_file, get_rank, get_partition_id
  from bigdl.orca.learn.log_monitor import LogMonitor
label_desc: Use image node_id 3 for calling the PytorchPysparkWorker obj's underlying member method code with example usage: obj.setup_distributed(mode, cluster_info, driver_ip, driver_tcp_store_port) without return types