code: def test_onnx_export_timm_image(checkpoint_name, num_gpus):
      model_path = "./automm_shopee"
      download_dir = "./ag_automm_tutorial_imgcls"
      train_data, test_data = shopee_dataset(download_dir)
      image_path_export = test_data.iloc[0]["image"]
      image_path_test = test_data.iloc[1]["image"]
  
      if os.path.exists(model_path):
          shutil.rmtree(model_path)
  
      # train
      predictor = MultiModalPredictor(
          hyperparameters={
              "optimization.max_epochs": 1,
              "model.names": ["timm_image"],
              "model.timm_image.checkpoint_name": checkpoint_name,
              "env.num_gpus": num_gpus,
              "env.num_workers": 0,
              "env.strategy": "ddp",
          },
          label="label",
          path=model_path,
      )
      predictor.fit(
          train_data=train_data,
          time_limit=10,  # seconds
      )
      predictor.save(path=model_path)
      loaded_predictor = MultiModalPredictor.load(path=model_path)
  
      # predict
      load_proba = loaded_predictor.predict_proba(
          {"image": [image_path_test]}
      )
  
      # -------------------------------------------------------
      # convert (no path)
      onnx_path = loaded_predictor.export_onnx(
          {"image": [image_path_export]}
      )
      assert isinstance(
          onnx_path, bytes
      ), f"export_onnx() method should return onnx bytes directly, if path is not provided."
  
      # create onnx module for evaluation
      onnx_module = OnnxModule(
          onnx_path, providers=["CUDAExecutionProvider"]
      )
      onnx_module.input_keys = (
          loaded_predictor._learner._model.input_keys
      )
      onnx_module.prefix = loaded_predictor._learner._model.prefix
      onnx_module.get_output_dict = (
          loaded_predictor._learner._model.get_output_dict
      )
  
      # simply replace _model in the loaded predictor to predict with onnxruntime
      loaded_predictor._learner._model = onnx_module
      onnx_proba = loaded_predictor.predict_proba(
          {"image": [image_path_test]}
      )
  
      # assert allclose
      np.testing.assert_allclose(
          load_proba, onnx_proba, rtol=1e-2, atol=1e-2
      )
  
      # -------------------------------------------------------
      # convert (with path)
      loaded_predictor = MultiModalPredictor.load(path=model_path)
      onnx_path = loaded_predictor.export_onnx(
          {"image": [image_path_export]}, path=model_path
      )
  
      # Check existence of the exported onnx model file and tensorrt cache files
      onnx_path_expected = os.path.join(model_path, "model.onnx")
      assert isinstance(
          onnx_path, str
      ), "export_onnx() method should return a string, if path argument is provided."
      assert onnx_path == onnx_path_expected, "onnx_path mismatch"
      assert os.path.exists(
          onnx_path
      ), f"onnx model file not found at {onnx_path}"
  
      # create onnx module for evaluation
      onnx_module = OnnxModule(
          onnx_path, providers=["CUDAExecutionProvider"]
      )
      onnx_module.input_keys = (
          loaded_predictor._learner._model.input_keys
      )
      onnx_module.prefix = loaded_predictor._learner._model.prefix
      onnx_module.get_output_dict = (
          loaded_predictor._learner._model.get_output_dict
      )
  
      # simply replace _model in the loaded predictor to predict with onnxruntime
      loaded_predictor._learner._model = onnx_module
      onnx_proba = loaded_predictor.predict_proba(
          {"image": [image_path_test]}
      )
  
      # assert allclose
      np.testing.assert_allclose(
          load_proba, onnx_proba, rtol=1e-2, atol=1e-2
      )
  
file_path: autogluon/multimodal/tests/unittests/others/test_deployment_onnx.py
imports: import os
  import shutil
  import numpy
  import numpy.testing
  import pytest
  from packaging import version
  from scipy.stats import pearsonr, spearmanr
  from sklearn.metrics.pairwise import paired_cosine_distances
  from autogluon.multimodal import MultiModalPredictor
  from autogluon.multimodal.constants import REGRESSION
  from autogluon.multimodal.utils.misc import shopee_dataset
  from autogluon.multimodal.utils.onnx import OnnxModule
  from ..utils.unittest_datasets import AEDataset, PetFinderDataset
label_desc: Use image node_id 3 for calling a global function with example usage: test_onnx_export_timm_image(checkpoint_name, num_gpus) without return types