code: def model_fn(model_dir):
      logger.info("model_fn")
      neopytorch.config(model_dir=model_dir, neo_runtime=True)
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
      # The compiled model is saved as "model.pth" or "model.pt"
      model = torch.jit.load(
          os.path.join(model_dir, "model.pt"), map_location=device
      )
  
      # It is recommended to run warm-up inference during model load
      sample_input_path = os.path.join(model_dir, "sample_input.pkl")
      with open(sample_input_path, "rb") as input_file:
          model_input = pickle.load(input_file)
      if torch.is_tensor(model_input):
          model_input = model_input.to(device)
          model(model_input)
      elif isinstance(model_input, tuple):
          model_input = (
              inp.to(device)
              for inp in model_input
              if torch.is_tensor(inp)
          )
          model(*model_input)
      else:
          print(
              "Only supports a torch tensor or a tuple of torch tensors"
          )
  
      return model
  
file_path: sagemaker-python-sdk/tests/data/pytorch_neo/code/inference.py
imports: import io
  import json
  import logging
  import os
  import pickle
  import numpy
  import torch
  import neopytorch
  import torchvision.transforms
  from PIL import Image
label_desc: Use image node_id 2 for calling a global function with example usage: model_fn(model_dir) and returns: model