code: def test_onnx_config_with_loss(self):
      # Prepare model and dataset
      model_checkpoint = "hf-internal-testing/tiny-random-bert"
      models = {
          AutoModelForSequenceClassification.from_pretrained(
              model_checkpoint
          ),
          TFAutoModelForSequenceClassification.from_pretrained(
              model_checkpoint
          ),
      }
  
      for model in models:
          with self.subTest(model=model):
              with tempfile.TemporaryDirectory() as tmp_dir:
                  onnx_config_constructor = (
                      TasksManager.get_exporter_config_constructor(
                          model=model,
                          exporter="onnx",
                          task="text-classification",
                      )
                  )
                  onnx_config = onnx_config_constructor(model.config)
  
                  wrapped_onnx_config = OnnxConfigWithLoss(onnx_config)
  
                  # Export model from PyTorch to ONNX
                  onnx_model_path = Path(
                      os.path.join(
                          tmp_dir, f"{model.config.model_type}.onnx"
                      )
                  )
                  opset = max(onnx_config.DEFAULT_ONNX_OPSET, 12)
                  _ = export(
                      model=model,
                      config=wrapped_onnx_config,
                      opset=opset,
                      output=onnx_model_path,
                  )
  
                  # ONNX Runtime Inference
                  ort_sess = onnxruntime.InferenceSession(
                      onnx_model_path.as_posix(),
                      providers=[
                          "CUDAExecutionProvider"
                          if torch.cuda.is_available()
                          and "CUDAExecutionProvider"
                          in onnxruntime.get_available_providers()
                          else "CPUExecutionProvider"
                      ],
                  )
                  framework = (
                      "pt"
                      if isinstance(model, PreTrainedModel)
                      else "tf"
                  )
                  normalized_config = NormalizedConfigManager.get_normalized_config_class(
                      "bert"
                  )(
                      model.config
                  )
                  input_generator = DummyTextInputGenerator(
                      "text-classification",
                      normalized_config,
                      batch_size=2,
                      sequence_length=16,
                  )
  
                  inputs = {
                      name: input_generator.generate(
                          name, framework=framework
                      )
                      for name in [
                          "input_ids",
                          "attention_mask",
                          "token_type_ids",
                      ]
                  }
                  inputs["labels"] = input_generator.constant_tensor(
                      [2],
                      value=0,
                      dtype=inputs["input_ids"].dtype,
                      framework=framework,
                  )
  
                  input_names = [
                      ort_input.name
                      for ort_input in ort_sess._inputs_meta
                  ]
                  output_names = [
                      output.name for output in ort_sess._outputs_meta
                  ]
  
                  input_feed = {
                      input_name: inputs[input_name].cpu().numpy()
                      for input_name in input_names
                  }
  
                  ort_outputs = ort_sess.run(output_names, input_feed)
                  pt_outputs = model(**inputs)
  
                  # Checkers
                  assert (
                      len(ort_outputs) > 1
                  ), "There is only one element in outputs, the loss might be missing!"
                  if issubclass(type(model), PreTrainedModel):
                      self.assertAlmostEqual(
                          float(ort_outputs[0]),
                          float(pt_outputs["loss"]),
                          3,
                          "The losses of ONNX Runtime and PyTorch inference are not close enough!",
                      )
                  elif issubclass(type(model), TFPreTrainedModel):
                      for ort_loss, pt_loss in zip(
                          ort_outputs[-1], pt_outputs["loss"]
                      ):
                          self.assertAlmostEqual(
                              float(ort_loss),
                              float(pt_loss),
                              3,
                              "The losses of ONNX Runtime and PyTorch inference are not close enough!",
                          )
                  gc.collect()
  
imports: import gc
  import os
  import tempfile
  import unittest
  from pathlib import Path
  import onnxruntime
  import pytest
  import torch
  from transformers import AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
  from transformers.modeling_utils import PreTrainedModel
  from transformers.utils import is_tf_available
  from optimum.exporters import TasksManager
  from optimum.exporters.onnx import OnnxConfigWithLoss, export
  from optimum.onnxruntime.utils import ONNX_DECODER_NAME
  from optimum.utils import DummyTextInputGenerator
  from optimum.utils.normalized_config import NormalizedConfigManager
label_desc: Use image node_id 1 for calling the TestOnnxConfigWithLoss obj's underlying member method code with example usage: obj.test_onnx_config_with_loss() without return types