code: def test_models_onnx_export(arch_name, input_shape, output_size):
      # Model
      batch_size = 2
      tf.keras.backend.clear_session()
      if arch_name == "mobilenet_v3_small_orientation":
          model = classification.__dict__[arch_name](
              pretrained=True, input_shape=input_shape
          )
      else:
          model = classification.__dict__[arch_name](
              pretrained=True, include_top=True, input_shape=input_shape
          )
  
      if arch_name == "vit_b":
          # vit model needs a fixed batch size
          dummy_input = [
              tf.TensorSpec([2, *input_shape], tf.float32, name="input")
          ]
      else:
          # batch_size = None for dynamic batch size
          dummy_input = [
              tf.TensorSpec(
                  [None, *input_shape], tf.float32, name="input"
              )
          ]
  
      np_dummy_input = np.random.rand(batch_size, *input_shape).astype(
          np.float32
      )
      with tempfile.TemporaryDirectory() as tmpdir:
          # Export
          model_path, output = export_model_to_onnx(
              model,
              model_name=os.path.join(tmpdir, "model"),
              dummy_input=dummy_input,
          )
  
          assert os.path.exists(model_path)
          # Inference
          ort_session = onnxruntime.InferenceSession(
              os.path.join(tmpdir, "model.onnx"),
              providers=["CPUExecutionProvider"],
          )
          ort_outs = ort_session.run(output, {"input": np_dummy_input})
          assert isinstance(ort_outs, list) and len(ort_outs) == 1
          assert ort_outs[0].shape == (batch_size, *output_size)
  
file_path: doctr/tests/tensorflow/test_models_classification_tf.py
imports: import os
  import tempfile
  import cv2
  import numpy
  import onnxruntime
  import psutil
  import pytest
  import tensorflow
  from doctr.models import classification
  from doctr.models.classification.predictor import CropOrientationPredictor
  from doctr.models.utils import export_model_to_onnx
label_desc: Use image node_id 5 for calling a global function with example usage: test_models_onnx_export(arch_name, input_shape, output_size) without return types