code: def test_triton_torchscript(csv_filename, tmpdir):
      # Configure features to be tested:
      input_features = [
          binary_feature(),
          number_feature(),
          category_feature(encoder={"vocab_size": 3}),
          # TODO: future support
          # sequence_feature(encoder={"vocab_size": 3}),
          # text_feature(encoder={"vocab_size": 3}),
          # vector_feature(),
          # timeseries_feature(),
          # date_feature(),
          # h3_feature(),
          # set_feature(encoder={"vocab_size": 3}),
          # bag_feature(encoder={"vocab_size": 3}),
          # image_feature(image_dest_folder),
          # audio_feature(audio_dest_folder),
      ]
      output_features = [
          binary_feature(),
          number_feature(),
          category_feature(decoder={"vocab_size": 3}),
          sequence_feature(decoder={"vocab_size": 3}),
          text_feature(decoder={"vocab_size": 3}),
          set_feature(decoder={"vocab_size": 3}),
          vector_feature(),
      ]
      backend = LocalTestBackend()
      config = {
          "input_features": input_features,
          "output_features": output_features,
          TRAINER: {"epochs": 1, BATCH_SIZE: 128},
      }
  
      # Generate training data
      training_data_csv_path = generate_data(
          input_features, output_features, csv_filename
      )
  
      # Train Ludwig (Pythonic) model:
      ludwig_model = LudwigModel(config, backend=backend)
      ludwig_model.train(
          dataset=training_data_csv_path,
          skip_save_training_description=True,
          skip_save_training_statistics=True,
          skip_save_model=True,
          skip_save_progress=True,
          skip_save_log=True,
          skip_save_processed_input=True,
      )
  
      # Create graph inference model (Torchscript) from trained Ludwig model.
      triton_path = os.path.join(tmpdir, "triton")
      model_name = "test_triton"
      model_version = "1"
      df = pd.read_csv(training_data_csv_path)
      triton_artifacts = export_triton(
          model=ludwig_model,
          data_example=df,
          model_name=model_name,
          output_path=triton_path,
          model_version=model_version,
      )
  
      # Validate that artifact paths exist.
      assert os.path.isdir(triton_path)
      assert all(
          os.path.exists(artifact.path) for artifact in triton_artifacts
      )
  
      # Load TorchScript models exported for Triton.
      triton_preprocessor = (
          triton_predictor
      ) = triton_postprocessor = None
      for artifact in triton_artifacts:
          if (
              artifact.model_name.endswith(PREPROCESSOR)
              and artifact.content_type == "application/octet-stream"
          ):
              triton_preprocessor = torch.jit.load(artifact.path)
          if (
              artifact.model_name.endswith(PREDICTOR)
              and artifact.content_type == "application/octet-stream"
          ):
              triton_predictor = torch.jit.load(artifact.path)
          if (
              artifact.model_name.endswith(POSTPROCESSOR)
              and artifact.content_type == "application/octet-stream"
          ):
              triton_postprocessor = torch.jit.load(artifact.path)
  
      assert triton_preprocessor is not None
      assert triton_predictor is not None
      assert triton_postprocessor is not None
  
      # Forward data through models.
      data_to_predict = to_inference_module_input_from_dataframe(
          df, ludwig_model.config, load_paths=True, device="cpu"
      )
      triton_preprocessor_output = triton_preprocessor(
          *data_to_predict.values()
      )
      triton_predictor_output = triton_predictor(
          *triton_preprocessor_output
      )
      triton_postprocessor_output = triton_postprocessor(
          *triton_predictor_output
      )
  
      # Get TorchScript inference modules and forward data.
      inference_modules = get_inference_modules(ludwig_model, "cpu")
      preprocessor_output = inference_modules[0](data_to_predict)
      predictor_output = inference_modules[1](preprocessor_output)
      postprocessor_output = inference_modules[2](predictor_output)
  
      assert len(postprocessor_output) == len(
          triton_postprocessor_output
      ), "Number of output mismatch after postprocessor step"
  
      for i, (_, out_value) in enumerate(postprocessor_output.items()):
          both_list = isinstance(out_value, list) and isinstance(
              triton_postprocessor_output[i], list
          )
          both_tensor = isinstance(
              out_value, torch.Tensor
          ) and isinstance(triton_postprocessor_output[i], torch.Tensor)
          assert (
              both_list or both_tensor
          ), "Type mismatch in PREDICTIONS, PROBABILITIES, LOGITS output"
  
          if (
              isinstance(out_value, list)
              and len(out_value) > 0
              and isinstance(out_value[0], str)
          ):
              assert (
                  out_value == triton_postprocessor_output[i]
              ), "Category feature outputs failure."
          elif (
              isinstance(out_value, list)
              and len(out_value) > 0
              and isinstance(out_value[0], torch.Tensor)
          ):
              assert len(out_value) == len(
                  triton_postprocessor_output[i]
              ), "Set feature outputs failure."
              assert all(
                  torch.allclose(inf, trit)
                  for inf, trit in zip(
                      out_value, triton_postprocessor_output[i]
                  )
              ), "Set feature outputs failure."
          elif (
              isinstance(out_value, list)
              and len(out_value) > 0
              and isinstance(out_value[0], list)
          ):
              assert len(out_value) == len(
                  triton_postprocessor_output[i]
              ), "Sequence (including text, etc.) feature outputs failure."
              assert all(
                  inf == trit
                  for inf, trit in zip(
                      out_value, triton_postprocessor_output[i]
                  )
              ), "Sequence (including text, etc.) feature outputs failure."
          elif isinstance(out_value, torch.Tensor):
              assert torch.allclose(
                  out_value, triton_postprocessor_output[i]
              )
          else:
              raise ValueError(
                  "Value should be either List[str] or torch.Tensor."
              )
  
file_path: ludwig/tests/integration_tests/test_triton.py
imports: import os
  from typing import List
  import pandas
  import pytest
  import torch
  from ludwig.api import LudwigModel
  from ludwig.constants import BATCH_SIZE, TRAINER
  from ludwig.data.dataset_synthesizer import build_synthetic_dataset_df
  from ludwig.utils.data_utils import load_yaml
  from ludwig.utils.inference_utils import to_inference_module_input_from_dataframe
  from ludwig.utils.triton_utils import export_triton, get_inference_modules, POSTPROCESSOR, PREDICTOR, PREPROCESSOR
  from tests.integration_tests.utils import binary_feature, category_feature, generate_data, LocalTestBackend, number_feature, sequence_feature, set_feature, text_feature, vector_feature
label_desc: Use image node_id 1 for calling a global function with example usage: test_triton_torchscript(csv_filename, tmpdir) without return types