code: def test_classification_model(model_fn, dev):
      set_rng_seed(0)
      defaults = {
          "num_classes": 50,
          "input_shape": (1, 3, 224, 224),
      }
      model_name = model_fn.__name__
      if SKIP_BIG_MODEL and is_skippable(model_name, dev):
          pytest.skip(
              "Skipped to reduce memory usage. Set env var SKIP_BIG_MODEL=0 to enable test for this model"
          )
      kwargs = {**defaults, **_model_params.get(model_name, {})}
      num_classes = kwargs.get("num_classes")
      input_shape = kwargs.pop("input_shape")
      real_image = kwargs.pop("real_image", False)
  
      model = model_fn(**kwargs)
      model.eval().to(device=dev)
      x = _get_image(
          input_shape=input_shape, real_image=real_image, device=dev
      )
      out = model(x)
      # FIXME: this if/else is nasty and only here to please our CI prior to the
      # release. We rethink these tests altogether.
      if model_name == "resnet101":
          prec = 0.2
      else:
          # FIXME: this is probably still way too high.
          prec = 0.1
      _assert_expected(out.cpu(), model_name, prec=prec)
      assert out.shape[-1] == num_classes
      _check_jit_scriptable(
          model,
          (x,),
          unwrapper=script_model_unwrapper.get(model_name, None),
          eager_out=out,
      )
      _check_fx_compatible(model, x, eager_out=out)
  
      if dev == "cuda":
          with torch.cuda.amp.autocast():
              out = model(x)
              # See autocast_flaky_numerics comment at top of file.
              if model_name not in autocast_flaky_numerics:
                  _assert_expected(out.cpu(), model_name, prec=0.1)
              assert out.shape[-1] == 50
  
      _check_input_backprop(model, x)
  
imports: import contextlib
  import functools
  import operator
  import os
  import pkgutil
  import platform
  import sys
  import warnings
  from collections import OrderedDict
  from tempfile import TemporaryDirectory
  from typing import Any
  import pytest
  import torch
  import torch.fx
  import torch.nn
  from _utils_internal import get_relative_path
  from common_utils import cpu_and_cuda, freeze_rng_state, map_nested_tensor_object, needs_cuda, set_rng_seed
  from PIL import Image
  from torchvision import models, transforms
  from torchvision.models import get_model_builder, list_models
label_desc: Use image node_id 22 for calling a global function with example usage: test_classification_model(model_fn, dev) without return types