code: def test_tensorrt_torch(
      output_library: DeepLearningFramework,
      dynamic: bool,
      quantization_type: QuantizationType,
      metric_drop_ths: int,
      metric: str,
  ):
      with TemporaryDirectory() as tmp_dir:
          (
              model,
              input_data,
              model_params,
              input_tfms,
              model_outputs,
              metric,
          ) = initialize_model(dynamic, metric, output_library, device)
  
          compiler_op = PyTorchTensorRTCompiler()
          compiler_op.to(device).execute(
              model=model,
              model_params=model_params,
              input_tfms=input_tfms,
              metric_drop_ths=metric_drop_ths,
              quantization_type=quantization_type,
              input_data=input_data,
          )
  
          compiled_model = compiler_op.get_result()
  
          build_inference_learner_op = (
              COMPILER_TO_INFERENCE_LEARNER_MAP[
                  ModelCompiler.TENSOR_RT_TORCH
              ]()
          )
  
          build_inference_learner_op.to(device).execute(
              model=compiled_model,
              model_orig=compiler_op.model_orig
              if hasattr(compiler_op, "model_orig")
              else None,
              model_params=model_params,
              input_tfms=input_tfms,
              source_dl_framework=output_library,
          )
  
          optimized_model = build_inference_learner_op.get_result()
          assert isinstance(
              optimized_model, PytorchTensorRTInferenceLearner
          )
  
          # Test save and load functions
          optimized_model.save(tmp_dir)
          loaded_model = PytorchTensorRTInferenceLearner.load(tmp_dir)
          assert isinstance(
              loaded_model, PytorchTensorRTInferenceLearner
          )
  
          assert isinstance(optimized_model.get_size(), int)
  
          inputs_example = tuple(optimized_model.get_inputs_example())
          res = optimized_model(*inputs_example)
          assert res is not None
  
          res_loaded = loaded_model(*inputs_example)
          assert all(
              [
                  torch.allclose(res_tensor, res_loaded_tensor)
                  for (res_tensor, res_loaded_tensor) in zip(
                      res, res_loaded
                  )
              ]
          )
  
          # Test validity of the model
          valid = check_model_validity(
              optimized_model,
              input_data,
              model_outputs,
              metric_drop_ths,
              quantization_type,
              metric,
          )
          assert valid
  
          if dynamic:  # Check also with a smaller bath_size
              torch_device = torch.device(
                  "cuda" if torch.cuda.is_available() else "cpu"
              )
  
              inputs_example = [
                  input_[: len(input_) // 2].to(torch_device)
                  for input_ in inputs_example
              ]
              res = optimized_model(*inputs_example)
              assert res is not None
  
              res_orig = tuple(model(*inputs_example))
              assert all(
                  [
                      torch.allclose(
                          res_tensor.float(),
                          res_orig_tensor,
                          rtol=1e-01,
                      )
                      for (res_tensor, res_orig_tensor) in zip(
                          res, res_orig
                      )
                  ]
              )
  
file_path: nebuly/optimization/nebullvm/nebullvm/operations/optimizations/tests/test_tensor_rt.py
imports: from pathlib import Path
  from tempfile import TemporaryDirectory
  import pytest
  import torch
  from nebullvm.core.models import Device, DeviceType, DeepLearningFramework, QuantizationType, ModelCompiler
  from nebullvm.operations.conversions.converters import PytorchConverter
  from nebullvm.operations.inference_learners.tensor_rt import TENSOR_RT_INFERENCE_LEARNERS, PytorchTensorRTInferenceLearner
  from nebullvm.operations.optimizations.compilers.tensor_rt import ONNXTensorRTCompiler, PyTorchTensorRTCompiler
  from nebullvm.operations.optimizations.optimizers.base import COMPILER_TO_INFERENCE_LEARNER_MAP
  from nebullvm.operations.optimizations.tests.utils import initialize_model, check_model_validity
  from nebullvm.operations.inference_learners.utils import load_model
  from nebullvm.tools.utils import check_module_version
label_desc: Use image node_id 2 for calling a global function with example usage: test_tensorrt_torch(output_library, dynamic, quantization_type, metric_drop_ths, metric) without return types