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","."]
docstr: Runs inference on a single batch of tensor data.
  
  This method is called by `TorchPredictor.predict` after converting the
  original data batch to torch tensors.
  
  Override this method to add custom logic for processing the model input or
  output.
  
  Args:
      inputs: A batch of data to predict on, represented as either a single
          PyTorch tensor or for multi-input models, a dictionary of tensors.
  
  Returns:
      The model outputs, either as a single tensor or a dictionary of tensors.
  
  Example:
  
      .. testcode::
  
          import numpy as np
          import torch
          from ray.train.torch import TorchPredictor
  
          # List outputs are not supported by default TorchPredictor.
          # So let's define a custom TorchPredictor and override call_model
          class MyModel(torch.nn.Module):
              def forward(self, input_tensor):
                  return [input_tensor, input_tensor]
  
          # Use a custom predictor to format model output as a dict.
          class CustomPredictor(TorchPredictor):
              def call_model(self, inputs):
                  model_output = super().call_model(inputs)
                  return {
                      str(i): model_output[i] for i in range(len(model_output))
                  }
  
          # create our data batch
          data_batch = np.array([1, 2])
          # create custom predictor and predict
          predictor = CustomPredictor(model=MyModel())
          predictions = predictor.predict(data_batch)
          print(f"Predictions: {predictions.get('0')}, {predictions.get('1')}")
  
      .. testoutput::
  
          Predictions: [1 2], [1 2]
docstr_tok: ["Runs","inference","on","a","single","batch","of","tensor","data",".","This","method","is","called","by","`","TorchPredictor.predict","`","after","converting","the","original","data","batch","to","torch","tensors",".","Override","this","method","to","add","custom","logic","for","processing","the","model","input","or","output",".","Args",":","inputs",":","A","batch","of","data","to","predict","on",",","represented","as","either","a","single","PyTorch","tensor","or","for","multi-input","models",",","a","dictionary","of","tensors",".","Returns",":","The","model","outputs",",","either","as","a","single","tensor","or","a","dictionary","of","tensors",".","Example",":","..","testcode",":",":","import","numpy","as","np","import","torch","from","ray.train.torch","import","TorchPredictor","#","List","outputs","are","not","supported","by","default","TorchPredictor",".","#","So","let","'s","define","a","custom","TorchPredictor","and","override","call_model","class","MyModel","(","torch.nn.Module",")",":","def","forward","(","self",",","input_tensor",")",":","return","[","input_tensor",",","input_tensor","]","#","Use","a","custom","predictor","to","format","model","output","as","a","dict",".","class","CustomPredictor","(","TorchPredictor",")",":","def","call_model","(","self",",","inputs",")",":","model_output","=","super","(",")",".call_model","(","inputs",")","return","{","str","(","i",")",":","model_output","[","i","]","for","i","in","range","(","len","(","model_output",")",")","}","#","create","our","data","batch","data_batch","=","np.array","(","[","1",",","2","]",")","#","create","custom","predictor","and","predict","predictor","=","CustomPredictor","(","model=MyModel","(",")",")","predictions","=","predictor.predict","(","data_batch",")","print","(","f","''","Predictions",":","{","predictions.get","(","'","0","'",")","}",",","{","predictions.get","(","'","1","'",")","}","''",")","..","testoutput",":",":","Predictions",":","[","1","2","]",",","[","1","2","]"]
code: def call_model(
      self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]]
  ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
      """Runs inference on a single batch of tensor data.
  
      This method is called by `TorchPredictor.predict` after converting the
      original data batch to torch tensors.
  
      Override this method to add custom logic for processing the model input or
      output.
  
      Args:
          inputs: A batch of data to predict on, represented as either a single
              PyTorch tensor or for multi-input models, a dictionary of tensors.
  
      Returns:
          The model outputs, either as a single tensor or a dictionary of tensors.
  
      Example:
  
          .. testcode::
  
              import numpy as np
              import torch
              from ray.train.torch import TorchPredictor
  
              # List outputs are not supported by default TorchPredictor.
              # So let's define a custom TorchPredictor and override call_model
              class MyModel(torch.nn.Module):
                  def forward(self, input_tensor):
                      return [input_tensor, input_tensor]
  
              # Use a custom predictor to format model output as a dict.
              class CustomPredictor(TorchPredictor):
                  def call_model(self, inputs):
                      model_output = super().call_model(inputs)
                      return {
                          str(i): model_output[i] for i in range(len(model_output))
                      }
  
              # create our data batch
              data_batch = np.array([1, 2])
              # create custom predictor and predict
              predictor = CustomPredictor(model=MyModel())
              predictions = predictor.predict(data_batch)
              print(f"Predictions: {predictions.get('0')}, {predictions.get('1')}")
  
          .. testoutput::
  
              Predictions: [1 2], [1 2]
  
      """
      with torch.no_grad():
          output = self.model(inputs)
      return output
  
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 4 for calling the TorchPredictor obj's underlying member method code with example usage: obj.call_model(inputs) and returns: output