code: from typing import Any, Type, Union, List
  
  import torch
  import torch.nn as nn
  from torch import Tensor
  from torch.quantization import fuse_modules
  from torchvision.models.resnet import Bottleneck, BasicBlock, ResNet, model_urls
  
  from ..._internally_replaced_utils import load_state_dict_from_url
  from .utils import _replace_relu, quantize_model
  
  __all__ = ["QuantizableResNet", "resnet18", "resnet50", "resnext101_32x8d"]
  
  
  quant_model_urls = {
      "resnet18_fbgemm": "https://download.pytorch.org/models/quantized/resnet18_fbgemm_16fa66dd.pth",
      "resnet50_fbgemm": "https://download.pytorch.org/models/quantized/resnet50_fbgemm_bf931d71.pth",
      "resnext101_32x8d_fbgemm": "https://download.pytorch.org/models/quantized/resnext101_32x8_fbgemm_09835ccf.pth",
  }
  
  
  class QuantizableBasicBlock(BasicBlock):
      def __init__(self, *args: Any, **kwargs: Any) -> None:
          super(QuantizableBasicBlock, self).__init__(*args, **kwargs)
          self.add_relu = torch.nn.quantized.FloatFunctional()
  
      def forward(self, x: Tensor) -> Tensor:
          identity = x
  
          out = self.conv1(x)
          out = self.bn1(out)
          out = self.relu(out)
  
          out = self.conv2(out)
          out = self.bn2(out)
  
          if self.downsample is not None:
              identity = self.downsample(x)
  
          out = self.add_relu.add_relu(out, identity)
  
          return out
  
      def fuse_model(self) -> None:
          torch.quantization.fuse_modules(self, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True)
          if self.downsample:
              torch.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True)
  
  
  class QuantizableBottleneck(Bottleneck):
      def __init__(self, *args: Any, **kwargs: Any) -> None:
          super(QuantizableBottleneck, self).__init__(*args, **kwargs)
          self.skip_add_relu = nn.quantized.FloatFunctional()
          self.relu1 = nn.ReLU(inplace=False)
          self.relu2 = nn.ReLU(inplace=False)
  
      def forward(self, x: Tensor) -> Tensor:
          identity = x
          out = self.conv1(x)
          out = self.bn1(out)
          out = self.relu1(out)
          out = self.conv2(out)
          out = self.bn2(out)
          out = self.relu2(out)
  
          out = self.conv3(out)
          out = self.bn3(out)
  
          if self.downsample is not None:
              identity = self.downsample(x)
          out = self.skip_add_relu.add_relu(out, identity)
  
          return out
  
      def fuse_model(self) -> None:
          fuse_modules(self, [["conv1", "bn1", "relu1"], ["conv2", "bn2", "relu2"], ["conv3", "bn3"]], inplace=True)
          if self.downsample:
              torch.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True)
  
  
  class QuantizableResNet(ResNet):
      def __init__(self, *args: Any, **kwargs: Any) -> None:
          super(QuantizableResNet, self).__init__(*args, **kwargs)
  
          self.quant = torch.quantization.QuantStub()
          self.dequant = torch.quantization.DeQuantStub()
  
      def forward(self, x: Tensor) -> Tensor:
          x = self.quant(x)
          # Ensure scriptability
          # super(QuantizableResNet,self).forward(x)
          # is not scriptable
          x = self._forward_impl(x)
          x = self.dequant(x)
          return x
  
      def fuse_model(self) -> None:
          r"""Fuse conv/bn/relu modules in resnet models
  
          Fuse conv+bn+relu/ Conv+relu/conv+Bn modules to prepare for quantization.
          Model is modified in place.  Note that this operation does not change numerics
          and the model after modification is in floating point
          """
  
          fuse_modules(self, ["conv1", "bn1", "relu"], inplace=True)
          for m in self.modules():
              if type(m) == QuantizableBottleneck or type(m) == QuantizableBasicBlock:
                  m.fuse_model()
  
  
  def _resnet(
      arch: str,
      block: Type[Union[BasicBlock, Bottleneck]],
      layers: List[int],
      pretrained: bool,
      progress: bool,
      quantize: bool,
      **kwargs: Any,
  ) -> QuantizableResNet:
  
      model = QuantizableResNet(block, layers, **kwargs)
      _replace_relu(model)
      if quantize:
          # TODO use pretrained as a string to specify the backend
          backend = "fbgemm"
          quantize_model(model, backend)
      else:
          assert pretrained in [True, False]
  
      if pretrained:
          if quantize:
              model_url = quant_model_urls[arch + "_" + backend]
          else:
              model_url = model_urls[arch]
  
          state_dict = load_state_dict_from_url(model_url, progress=progress)
  
          model.load_state_dict(state_dict)
      return model
  
  
  def resnet18(
      pretrained: bool = False,
      progress: bool = True,
      quantize: bool = False,
      **kwargs: Any,
  ) -> QuantizableResNet:
      r"""ResNet-18 model from
      `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
  
      Args:
          pretrained (bool): If True, returns a model pre-trained on ImageNet
          progress (bool): If True, displays a progress bar of the download to stderr
          quantize (bool): If True, return a quantized version of the model
      """
      return _resnet("resnet18", QuantizableBasicBlock, [2, 2, 2, 2], pretrained, progress, quantize, **kwargs)
  
  
  def resnet50(
      pretrained: bool = False,
      progress: bool = True,
      quantize: bool = False,
      **kwargs: Any,
  ) -> QuantizableResNet:
  
      r"""ResNet-50 model from
      `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
  
      Args:
          pretrained (bool): If True, returns a model pre-trained on ImageNet
          progress (bool): If True, displays a progress bar of the download to stderr
          quantize (bool): If True, return a quantized version of the model
      """
      return _resnet("resnet50", QuantizableBottleneck, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
  
  
  def resnext101_32x8d(
      pretrained: bool = False,
      progress: bool = True,
      quantize: bool = False,
      **kwargs: Any,
  ) -> QuantizableResNet:
      r"""ResNeXt-101 32x8d model from
      `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
  
      Args:
          pretrained (bool): If True, returns a model pre-trained on ImageNet
          progress (bool): If True, displays a progress bar of the download to stderr
          quantize (bool): If True, return a quantized version of the model
      """
      kwargs["groups"] = 32
      kwargs["width_per_group"] = 8
      return _resnet("resnext101_32x8d", QuantizableBottleneck, [3, 4, 23, 3], pretrained, progress, quantize, **kwargs)
   import contextlib
  import functools
  import os
  import random
  import shutil
  import tempfile
  
  import numpy as np
  import torch
  from PIL import Image
  from torchvision import io
  
  import __main__  # noqa: 401
  
  
  IN_CIRCLE_CI = os.getenv("CIRCLECI", False) == "true"
  IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None
  IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1"
  CUDA_NOT_AVAILABLE_MSG = "CUDA device not available"
  CIRCLECI_GPU_NO_CUDA_MSG = "We're in a CircleCI GPU machine, and this test doesn't need cuda."
  
  
  @contextlib.contextmanager
  def get_tmp_dir(src=None, **kwargs):
      tmp_dir = tempfile.mkdtemp(**kwargs)
      if src is not None:
          os.rmdir(tmp_dir)
          shutil.copytree(src, tmp_dir)
      try:
          yield tmp_dir
      finally:
          shutil.rmtree(tmp_dir)
  
  
  def set_rng_seed(seed):
      torch.manual_seed(seed)
      random.seed(seed)
  
  
  class MapNestedTensorObjectImpl(object):
      def __init__(self, tensor_map_fn):
          self.tensor_map_fn = tensor_map_fn
  
      def __call__(self, object):
          if isinstance(object, torch.Tensor):
              return self.tensor_map_fn(object)
  
          elif isinstance(object, dict):
              mapped_dict = {}
              for key, value in object.items():
                  mapped_dict[self(key)] = self(value)
              return mapped_dict
  
          elif isinstance(object, (list, tuple)):
              mapped_iter = []
              for iter in object:
                  mapped_iter.append(self(iter))
              return mapped_iter if not isinstance(object, tuple) else tuple(mapped_iter)
  
          else:
              return object
  
  
  def map_nested_tensor_object(object, tensor_map_fn):
      impl = MapNestedTensorObjectImpl(tensor_map_fn)
      return impl(object)
  
  
  def is_iterable(obj):
      try:
          iter(obj)
          return True
      except TypeError:
          return False
  
  
  @contextlib.contextmanager
  def freeze_rng_state():
      rng_state = torch.get_rng_state()
      if torch.cuda.is_available():
          cuda_rng_state = torch.cuda.get_rng_state()
      yield
      if torch.cuda.is_available():
          torch.cuda.set_rng_state(cuda_rng_state)
      torch.set_rng_state(rng_state)
  
  
  def cycle_over(objs):
      for idx, obj1 in enumerate(objs):
          for obj2 in objs[:idx] + objs[idx + 1 :]:
              yield obj1, obj2
  
  
  def int_dtypes():
      return (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
  
  
  def float_dtypes():
      return (torch.float32, torch.float64)
  
  
  @contextlib.contextmanager
  def disable_console_output():
      with contextlib.ExitStack() as stack, open(os.devnull, "w") as devnull:
          stack.enter_context(contextlib.redirect_stdout(devnull))
          stack.enter_context(contextlib.redirect_stderr(devnull))
          yield
  
  
  def cpu_and_gpu():
      import pytest  # noqa
  
      return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
  
  
  def needs_cuda(test_func):
      import pytest  # noqa
  
      return pytest.mark.needs_cuda(test_func)
  
  
  def _create_data(height=3, width=3, channels=3, device="cpu"):
      # TODO: When all relevant tests are ported to pytest, turn this into a module-level fixture
      tensor = torch.randint(0, 256, (channels, height, width), dtype=torch.uint8, device=device)
      data = tensor.permute(1, 2, 0).contiguous().cpu().numpy()
      mode = "RGB"
      if channels == 1:
          mode = "L"
          data = data[..., 0]
      pil_img = Image.fromarray(data, mode=mode)
      return tensor, pil_img
  
  
  def _create_data_batch(height=3, width=3, channels=3, num_samples=4, device="cpu"):
      # TODO: When all relevant tests are ported to pytest, turn this into a module-level fixture
      batch_tensor = torch.randint(0, 256, (num_samples, channels, height, width), dtype=torch.uint8, device=device)
      return batch_tensor
  
  
  assert_equal = functools.partial(torch.testing.assert_close, rtol=0, atol=1e-6)
  
  
  def get_list_of_videos(tmpdir, num_videos=5, sizes=None, fps=None):
      names = []
      for i in range(num_videos):
          if sizes is None:
              size = 5 * (i + 1)
          else:
              size = sizes[i]
          if fps is None:
              f = 5
          else:
              f = fps[i]
          data = torch.randint(0, 256, (size, 300, 400, 3), dtype=torch.uint8)
          name = os.path.join(tmpdir, "{}.mp4".format(i))
          names.append(name)
          io.write_video(name, data, fps=f)
  
      return names
  
  
  def _assert_equal_tensor_to_pil(tensor, pil_image, msg=None):
      np_pil_image = np.array(pil_image)
      if np_pil_image.ndim == 2:
          np_pil_image = np_pil_image[:, :, None]
      pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1)))
      if msg is None:
          msg = "tensor:\n{} \ndid not equal PIL tensor:\n{}".format(tensor, pil_tensor)
      assert_equal(tensor.cpu(), pil_tensor, msg=msg)
  
  
  def _assert_approx_equal_tensor_to_pil(
      tensor, pil_image, tol=1e-5, msg=None, agg_method="mean", allowed_percentage_diff=None
  ):
      # TODO: we could just merge this into _assert_equal_tensor_to_pil
      np_pil_image = np.array(pil_image)
      if np_pil_image.ndim == 2:
          np_pil_image = np_pil_image[:, :, None]
      pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1))).to(tensor)
  
      if allowed_percentage_diff is not None:
          # Assert that less than a given %age of pixels are different
          assert (tensor != pil_tensor).to(torch.float).mean() <= allowed_percentage_diff
  
      # error value can be mean absolute error, max abs error
      # Convert to float to avoid underflow when computing absolute difference
      tensor = tensor.to(torch.float)
      pil_tensor = pil_tensor.to(torch.float)
      err = getattr(torch, agg_method)(torch.abs(tensor - pil_tensor)).item()
      assert err < tol
  
  
  def _test_fn_on_batch(batch_tensors, fn, scripted_fn_atol=1e-8, **fn_kwargs):
      transformed_batch = fn(batch_tensors, **fn_kwargs)
      for i in range(len(batch_tensors)):
          img_tensor = batch_tensors[i, ...]
          transformed_img = fn(img_tensor, **fn_kwargs)
          assert_equal(transformed_img, transformed_batch[i, ...])
  
      if scripted_fn_atol >= 0:
          scripted_fn = torch.jit.script(fn)
          # scriptable function test
          s_transformed_batch = scripted_fn(batch_tensors, **fn_kwargs)
          torch.testing.assert_close(transformed_batch, s_transformed_batch, rtol=1e-5, atol=scripted_fn_atol)
  
apis: torch.quantization.DeQuantStub torch.nn.quantized.FloatFunctional torch.nn.ReLU torch.quantization.fuse_modules torch.quantization.QuantStub
apis: torch.jit.script torch.get_rng_state torch.randint torch.cuda.set_rng_state torch.cuda.get_rng_state torch.manual_seed torch.cuda.is_available torch.abs torch.set_rng_state numpy.array torch.testing.assert_close