code: def run_image_prediction(
      model: torch.nn.Module, images: torch.Tensor
  ) -> torch.Tensor:
      model.eval()
      with torch.no_grad():
          return torch.exp(model(images)).argmax(dim=1)
  
file_path: ray/python/ray/train/tests/test_horovod_trainer.py
imports: import os
  import pytest
  import torch
  import torch.nn
  from torch.utils.data import DataLoader
  from torchvision import datasets
  from torchvision.transforms import transforms
  import ray
  from ray.train import ScalingConfig
  from ray.train.examples.horovod.horovod_pytorch_example import Net
  from ray.train.examples.horovod.horovod_pytorch_example import train_func
  from ray.train.horovod import HorovodTrainer
  from ray.train.torch import TorchPredictor
label_desc: Use image node_id 2 for calling a global function with example usage: run_image_prediction(model, images) and returns: torch