code: def test_detection_models(
      arch_name, input_shape, output_size, out_prob, train_mode
  ):
      batch_size = 2
      model = detection.__dict__[arch_name](pretrained=True)
      model = model.train() if train_mode else model.eval()
      assert isinstance(model, torch.nn.Module)
      input_tensor = torch.rand((batch_size, *input_shape))
      target = [
          {
              CLASS_NAME: np.array(
                  [[0.5, 0.5, 1, 1], [0.5, 0.5, 0.8, 0.8]],
                  dtype=np.float32,
              )
          },
          {
              CLASS_NAME: np.array(
                  [[0.5, 0.5, 1, 1], [0.5, 0.5, 0.8, 0.9]],
                  dtype=np.float32,
              )
          },
      ]
      if torch.cuda.is_available():
          model.cuda()
          input_tensor = input_tensor.cuda()
      out = model(
          input_tensor,
          target,
          return_model_output=True,
          return_preds=not train_mode,
      )
      assert isinstance(out, dict)
      assert len(out) == 3 if not train_mode else len(out) == 2
      # Check proba map
      assert out["out_map"].shape == (batch_size, *output_size)
      assert out["out_map"].dtype == torch.float32
      if out_prob:
          assert torch.all(
              (out["out_map"] >= 0) & (out["out_map"] <= 1)
          )
      # Check boxes
      if not train_mode:
          for boxes_dict in out["preds"]:
              for boxes in boxes_dict.values():
                  assert boxes.shape[1] == 5
                  assert np.all(boxes[:, :2] < boxes[:, 2:4])
                  assert np.all(boxes[:, :4] >= 0) and np.all(
                      boxes[:, :4] <= 1
                  )
      # Check loss
      assert isinstance(out["loss"], torch.Tensor)
      # Check the rotated case (same targets)
      target = [
          {
              CLASS_NAME: np.array(
                  [
                      [[0.5, 0.5], [1, 0.5], [1, 1], [0.5, 1]],
                      [[0.5, 0.5], [0.8, 0.5], [0.8, 0.8], [0.5, 0.8]],
                  ],
                  dtype=np.float32,
              )
          },
          {
              CLASS_NAME: np.array(
                  [
                      [[0.5, 0.5], [1, 0.5], [1, 1], [0.5, 1]],
                      [[0.5, 0.5], [0.8, 0.5], [0.8, 0.9], [0.5, 0.9]],
                  ],
                  dtype=np.float32,
              )
          },
      ]
      loss = model(input_tensor, target)["loss"]
      assert (
          isinstance(loss, torch.Tensor)
          and ((loss - out["loss"]).abs() / loss).item() < 1e-1
      )
  
file_path: doctr/tests/pytorch/test_models_detection_pt.py
imports: import os
  import tempfile
  import numpy
  import onnxruntime
  import pytest
  import torch
  from doctr.file_utils import CLASS_NAME
  from doctr.models import detection
  from doctr.models.detection._utils import dilate, erode
  from doctr.models.detection.predictor import DetectionPredictor
  from doctr.models.utils import export_model_to_onnx
label_desc: Use image node_id 1 for calling a global function with example usage: test_detection_models(arch_name, input_shape, output_size, out_prob, train_mode) without return types