code: def __init__(
      self,
      *,
      model_creator,
      optimizer_creator,
      loss_creator=None,
      metrics=None,
      scheduler_creator=None,
      training_operator_cls=TrainingOperator,
      initialization_hook=None,
      config=None,
      scheduler_step_freq="batch",
      use_tqdm=False,
      backend="torch_distributed",
      workers_per_node=1,
  ):
      # todo remove ray_ctx to run on workers
      ray_ctx = RayContext.get()
      if not (
          isinstance(model_creator, types.FunctionType)
          and isinstance(optimizer_creator, types.FunctionType)
      ):  # Torch model is also callable.
          raise ValueError(
              "Must provide a function for both model_creator and optimizer_creator"
          )
  
      self.model_creator = model_creator
      self.optimizer_creator = optimizer_creator
      self.loss_creator = loss_creator
      self.scheduler_creator = scheduler_creator
      self.training_operator_cls = training_operator_cls
      self.scheduler_step_freq = scheduler_step_freq
      self.use_tqdm = use_tqdm
  
      if not training_operator_cls and not loss_creator:
          raise ValueError(
              "If a loss_creator is not provided, you must "
              "provide a custom training operator."
          )
  
      self.initialization_hook = initialization_hook
      self.config = {} if config is None else config
      worker_config = self.config.copy()
      params = dict(
          model_creator=self.model_creator,
          optimizer_creator=self.optimizer_creator,
          loss_creator=self.loss_creator,
          scheduler_creator=self.scheduler_creator,
          training_operator_cls=self.training_operator_cls,
          scheduler_step_freq=self.scheduler_step_freq,
          use_tqdm=self.use_tqdm,
          config=worker_config,
          metrics=metrics,
      )
  
      if backend == "torch_distributed":
          cores_per_node = (
              ray_ctx.ray_node_cpu_cores // workers_per_node
          )
          num_nodes = ray_ctx.num_ray_nodes * workers_per_node
          RemoteRunner = ray.remote(num_cpus=cores_per_node)(
              TorchRunner
          )
          self.remote_workers = [
              RemoteRunner.remote(**params) for i in range(num_nodes)
          ]
          ray.get(
              [
                  worker.setup.remote(cores_per_node)
                  for i, worker in enumerate(self.remote_workers)
              ]
          )
  
          head_worker = self.remote_workers[0]
          address = ray.get(head_worker.setup_address.remote())
  
          logger.info(
              f"initializing pytorch process group on {address}"
          )
  
          ray.get(
              [
                  worker.setup_torch_distribute.remote(
                      address, i, num_nodes
                  )
                  for i, worker in enumerate(self.remote_workers)
              ]
          )
  
      elif backend == "horovod":
          from zoo.orca.learn.horovod.horovod_ray_runner import (
              HorovodRayRunner,
          )
  
          self.horovod_runner = HorovodRayRunner(
              ray_ctx,
              worker_cls=TorchRunner,
              worker_param=params,
              workers_per_node=workers_per_node,
          )
          self.remote_workers = self.horovod_runner.remote_workers
          cores_per_node = self.horovod_runner.cores_per_node
          ray.get(
              [
                  worker.setup.remote(cores_per_node)
                  for i, worker in enumerate(self.remote_workers)
              ]
          )
  
          ray.get(
              [
                  worker.setup_horovod.remote()
                  for i, worker in enumerate(self.remote_workers)
              ]
          )
      else:
          raise Exception(
              'Only "torch_distributed" and "horovod" are supported '
              "values of backend, but got {}".format(backend)
          )
      self.num_workers = len(self.remote_workers)
  
imports: import io
  import types
  import logging
  import numbers
  import torch
  import numpy
  from zoo.orca.data.ray_xshards import RayXShards
  from zoo.orca.learn.pytorch.training_operator import TrainingOperator
  from zoo.orca.learn.pytorch.torch_runner import TorchRunner
  from zoo.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, update_predict_xshards, process_xshards_of_pandas_dataframe
  from zoo.ray import RayContext
  import ray
  from ray.exceptions import RayActorError
label_desc: Use image node_id 1 to create a new PyTorchRayEstimator object  with example: obj = PyTorchRayEstimator()