docstr: Using this PyTorch model to make predictions on the data.
  
  :param data: An instance of SparkXShards, a Ray Dataset or a Spark DataFrame
  :param batch_size: Total batch size for all workers used for inference. Each worker's batch
         size would be this value divide the total number of workers. Default is 32.
  :param profile: Boolean. Whether to return time stats for the training procedure.
         Default is False.
  :param feature_cols: feature column names if data is a Spark DataFrame or Ray Dataset.
  :return: A SparkXShards or a list that contains the predictions with key "prediction"
         in each shard
docstr_tok: ["Using","this","PyTorch","model","to","make","predictions","on","the","data",".",":","param","data",":","An","instance","of","SparkXShards",",","a","Ray","Dataset","or","a","Spark","DataFrame",":","param","batch_size",":","Total","batch","size","for","all","workers","used","for","inference",".","Each","worker","'s","batch","size","would","be","this","value","divide","the","total","number","of","workers",".","Default","is","32.",":","param","profile",":","Boolean",".","Whether","to","return","time","stats","for","the","training","procedure",".","Default","is","False",".",":","param","feature_cols",":","feature","column","names","if","data","is","a","Spark","DataFrame","or","Ray","Dataset",".",":","return",":","A","SparkXShards","or","a","list","that","contains","the","predictions","with","key","``","prediction","''","in","each","shard"]
code: def predict(
      self,
      data: Union["SparkXShards", "SparkDataFrame"],
      batch_size: int = 32,
      feature_cols: Optional[List[str]] = None,
      profile: bool = False,
      callbacks: Optional[List["Callback"]] = None,
  ) -> Union["SparkXShards", "SparkDataFrame"]:
      """
      Using this PyTorch model to make predictions on the data.
  
      :param data: An instance of SparkXShards, a Ray Dataset or a Spark DataFrame
      :param batch_size: Total batch size for all workers used for inference. Each worker's batch
             size would be this value divide the total number of workers. Default is 32.
      :param profile: Boolean. Whether to return time stats for the training procedure.
             Default is False.
      :param feature_cols: feature column names if data is a Spark DataFrame or Ray Dataset.
      :return: A SparkXShards or a list that contains the predictions with key "prediction"
             in each shard
      """
      invalidInputError(
          isinstance(batch_size, int) and batch_size > 0,
          "batch_size should be a positive integer",
      )
      batch_size = (
          batch_size // self.num_workers
      )  # Local batch size for each worker
      if batch_size <= 0:
          batch_size = 1
      from bigdl.orca.data import SparkXShards
  
      callbacks = callbacks or []
      make_only_mainCallback(callbacks)
      if self.use_tqdm:
          callbacks.append(TqdmCallback())
  
      params = dict(
          batch_size=batch_size, profile=profile, callbacks=callbacks
      )
  
      from pyspark.sql import DataFrame
  
      if isinstance(data, DataFrame):
          xshards, _ = dataframe_to_xshards(
              data,
              validation_data=None,
              feature_cols=feature_cols,
              label_cols=None,
              mode="predict",
              shard_size=batch_size,
          )
          pred_shards = self._predict_spark_xshards(xshards, params)
          result = convert_predict_xshards_to_dataframe(
              data, pred_shards
          )
      elif isinstance(data, SparkXShards):
          xshards = data.to_lazy()
          if xshards._get_class_name() == "pandas.core.frame.DataFrame":
              xshards = process_xshards_of_pandas_dataframe(
                  xshards, feature_cols
              )
              pred_shards = self._predict_spark_xshards(xshards, params)
              result = add_predict_to_pd_xshards(data, pred_shards)
          else:
              pred_shards = self._predict_spark_xshards(xshards, params)
              result = update_predict_xshards(data, pred_shards)
      elif isinstance(data, ray.data.Dataset):
          shards = data.split(
              n=self.num_workers, locality_hints=self.remote_workers
          )
  
          def data_creator(config, batch_size):
              torch_datashard = shard.to_torch(
                  feature_columns=feature_cols, batch_size=batch_size
              )
              return torch_datashard
  
          remote_worker_stats = []
          for shard, worker in zip(shards, self.remote_workers):
              worker_stats = worker.predict.remote(
                  data_creator, batch_size, profile, callbacks
              )
              remote_worker_stats.append(worker_stats)
          result = ray.data.from_numpy(remote_worker_stats).map(
              lambda r: {"prediction_result": r["value"]}
          )
      else:
          invalidInputError(
              False,
              "Only xshards, Spark DataFrame or Ray Dataset"
              " is supported for predict",
          )
  
      return result
  
imports: import types
  import copy
  import logging
  import math
  from bigdl.orca.data.ray_xshards import RayXShards
  from bigdl.orca.learn.pytorch.pytorch_ray_worker import PytorchRayWorker
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, update_predict_xshards, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards
  from bigdl.orca.ray import OrcaRayContext
  from bigdl.orca.learn.pytorch.core.base_ray_estimator import BaseRayEstimator
  from bigdl.orca.learn.pytorch.utils import process_stats, check_for_failure
  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
  import ray
  from bigdl.dllib.utils.log4Error import invalidInputError
  from typing import TYPE_CHECKING, Union, Optional, Callable, Dict, List
label_desc: Use image node_id 3 for calling the PyTorchRayEstimator obj's underlying member method code with example usage: obj.predict(data, batch_size, feature_cols, profile, callbacks) and returns: result, torch_datashard