code: def _test_classification(
      model, input_shape, output_size, batch_size=2
  ):
      # Forward
      with torch.no_grad():
          out = model(
              torch.rand(
                  (batch_size, *input_shape), dtype=torch.float32
              )
          )
      # Output checks
      assert isinstance(out, torch.Tensor)
      assert out.dtype == torch.float32
      assert out.numpy().shape == (batch_size, *output_size)
      # Check FP16
      if torch.cuda.is_available():
          model = model.half().cuda()
          with torch.no_grad():
              out = model(
                  torch.rand(
                      (batch_size, *input_shape), dtype=torch.float16
                  ).cuda()
              )
          assert out.dtype == torch.float16
  
file_path: doctr/tests/pytorch/test_models_classification_pt.py
imports: import os
  import tempfile
  import cv2
  import numpy
  import onnxruntime
  import pytest
  import torch
  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 1 for calling a global function with example usage: _test_classification(model, input_shape, output_size, batch_size) without return types