code: def __init__(
      self,
      *,
      model_creator: Union[Callable[[Dict], "Module"], None] = None,
      optimizer_creator: Union[
          Callable[["Module", Dict], "Optimizer"], None
      ] = None,
      loss_creator: Union[
          "Loss", Callable[[Dict], "Loss"], None
      ] = None,
      metrics: Union["Metric", List["Metric"], None] = None,
      scheduler_creator: Optional[
          Callable[[Dict], "LRScheduler"]
      ] = None,
      config: Optional[Dict] = None,
      use_tqdm: bool = False,
      workers_per_node: int = 1,
      sync_stats: bool = True,
      log_level: int = logging.INFO,
      model_dir: Optional[str] = None,
      log_to_driver: bool = True,
  ):
      logging.basicConfig(
          level=log_level,
          format="[%(asctime)s] %(levelname)-8s %(message)s",
          datefmt="%Y-%m-%d %H:%M:%S",
      )
      self.logger = logging.getLogger(__name__)
      if config is not None and "batch_size" in config:
          invalidInputError(
              False,
              "Please do not specify batch_size in config. Input batch_size in the"
              " fit/evaluate/predict function of the estimator instead.",
          )
      self.config = {} if config is None else config
  
      sc = OrcaContext.get_spark_context()
  
      if model_creator and not isinstance(
          model_creator, types.FunctionType
      ):
          # Torch model is also callable.
          invalidInputError(
              False, "Must provide a function for model_creator"
          )
  
      self.model_dir = parse_model_dir(model_dir)
      self.use_tqdm = use_tqdm
  
      self.model_creator = model_creator
      self.optimizer_creator = optimizer_creator
  
      num_nodes, cores_per_node = get_node_and_core_number()
      self.num_workers = num_nodes * workers_per_node
      self.total_cores = num_nodes * cores_per_node
      self.cores_per_worker = cores_per_node // workers_per_node
  
      # over partition to cover tasks all over the cluster
      self.workerRDD = sc.parallelize(
          list(range(self.total_cores * 4)), self.total_cores * 4
      ).repartition(self.num_workers)
  
      self.ip = get_node_ip()
      self.log_port = find_free_port()
      is_local = sc.master.startswith("local")
      self.need_to_log_to_driver = (not is_local) and log_to_driver
      if self.need_to_log_to_driver:
          self.log_server_thread = start_log_server(
              self.ip, self.log_port
          )
      self.tcp_store_port = find_free_port()
  
      self.worker_init_params = dict(
          model_creator=self.model_creator,
          optimizer_creator=optimizer_creator,
          loss_creator=loss_creator,
          scheduler_creator=scheduler_creator,
          config=copy.copy(self.config),
          metrics=metrics,
          size=self.num_workers,
          cores_per_worker=self.cores_per_worker,
          sync_stats=sync_stats,
          log_level=log_level,
          model_dir=self.model_dir,
          log_to_driver=self.need_to_log_to_driver,
          driver_ip=self.ip,
          driver_log_port=self.log_port,
          driver_tcp_store_port=self.tcp_store_port,
      )
  
      local_init_params = self.worker_init_params.copy()
      local_init_params["log_to_driver"] = False
      self.driver_runner = PytorchPysparkWorker(
          mode="predict", **local_init_params
      )
  
      if self.model_creator:
          self.state_dict = self.driver_runner.get_state_dict()
  
imports: import types
  import torch
  import copy
  import os
  import shutil
  import tempfile
  import logging
  import math
  from pyspark.sql.dataframe import DataFrame
  from bigdl.dllib.utils.file_utils import is_local_path
  from bigdl.orca.learn.pytorch.pytorch_pyspark_worker import PytorchPysparkWorker
  from bigdl.orca.learn.pytorch.utils import process_stats
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, make_data_creator, update_predict_xshards, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards
  from bigdl.orca.data import SparkXShards
  from bigdl.orca import OrcaContext
  from bigdl.orca.learn.base_estimator import BaseEstimator
  from bigdl.orca.data.file import get_remote_file_to_local, put_local_file_to_remote
  from bigdl.dllib.utils.common import get_node_and_core_number
  from bigdl.orca.learn.log_monitor import start_log_server, stop_log_server
  from bigdl.orca.learn.pytorch.callbacks.maincallback import make_only_mainCallback
  from bigdl.orca.learn.pytorch.callbacks.tqdm import TqdmCallback, is_tqdm_exists
  from bigdl.orca.learn.pytorch.callbacks.maxsteps import MaxstepsCallback
  from bigdl.orca.learn.utils import find_free_port, find_ip_and_free_port
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.dllib.utils.log4Error import invalidInputError
  from typing import TYPE_CHECKING, Union, Optional, Callable, Dict, List
  from bigdl.orca.learn.pytorch.callbacks import Callback
label_desc: Use image node_id 1 to create a new PyTorchPySparkEstimator object from inherited base classes: BaseEstimator with example: obj = PyTorchPySparkEstimator()