class_docstr: A predictor for PyTorch models.
  
  Args:
      model: The torch module to use for predictions.
      preprocessor: A preprocessor used to transform data batches prior
          to prediction.
      use_gpu: If set, the model will be moved to GPU on instantiation and
          prediction happens on GPU.
class_docstr_tok: ["A","predictor","for","PyTorch","models",".","Args",":","model",":","The","torch","module","to","use","for","predictions",".","preprocessor",":","A","preprocessor","used","to","transform","data","batches","prior","to","prediction",".","use_gpu",":","If","set",",","the","model","will","be","moved","to","GPU","on","instantiation","and","prediction","happens","on","GPU","."]
code: def __init__(
      self,
      model: torch.nn.Module,
      preprocessor: Optional["Preprocessor"] = None,
      use_gpu: bool = False,
  ):
      self.model = model
      self.model.eval()
      self.use_gpu = use_gpu
  
      if use_gpu:
          # TODO (jiaodong): #26249 Use multiple GPU devices with sharded input
          self.device = torch.device("cuda")
      else:
          self.device = torch.device("cpu")
  
      # Ensure input tensor and model live on the same device
      self.model.to(self.device)
  
      if (
          not use_gpu
          and torch.cuda.device_count() > 0
          and log_once("torch_predictor_not_using_gpu")
      ):
          logger.warning(
              "You have `use_gpu` as False but there are "
              f"{torch.cuda.device_count()} GPUs detected on host where "
              "prediction will only use CPU. Please consider explicitly "
              "setting `TorchPredictor(use_gpu=True)` or "
              "`batch_predictor.predict(ds, num_gpus_per_worker=1)` to "
              "enable GPU prediction."
          )
  
      super().__init__(preprocessor)
  
file_path: ray/python/ray/train/torch/torch_predictor.py
imports: import logging
  from typing import TYPE_CHECKING, Dict, Optional, Union
  import numpy
  import torch
  from ray.air._internal.torch_utils import convert_ndarray_batch_to_torch_tensor_batch
  from ray.train._internal.dl_predictor import DLPredictor
  from ray.train.predictor import DataBatchType
  from ray.train.torch import TorchCheckpoint
  from ray.util import log_once
  from ray.util.annotations import DeveloperAPI, PublicAPI
label_desc: Use image node_id 1 to create a new TorchPredictor object from inherited base classes: DLPredictor with example: obj = TorchPredictor(model, preprocessor, use_gpu)