hunk.code_window: @pytest.mark.parametrize('batch_size', [1])
   def test_model_forward_torchscript(model_name, batch_size):
       """Run a single forward pass with each model"""
       with set_scriptable(True):
           model = create_model(model_name, pretrained=False)
       model.eval()
  
hunk.labels: keep keep add keep keep keep keep
hunk.after_edit:     input_size = _get_input_size(model_name=model_name, target=TARGET_JIT_SIZE)
       if max(input_size) > MAX_JIT_SIZE:  # NOTE using MAX_FWD_SIZE as the final limit is intentional
           pytest.skip("Fixed input size model > limit.")
  
hunk.file_path: tests/test_models.py
hunk.type: add
file: """ LeViT
  
  Paper: `LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference`
      - https://arxiv.org/abs/2104.01136
  
  @article{graham2021levit,
    title={LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference},
    author={Benjamin Graham and Alaaeldin El-Nouby and Hugo Touvron and Pierre Stock and Armand Joulin and Herv\'e J\'egou and Matthijs Douze},
    journal={arXiv preprint arXiv:22104.01136},
    year={2021}
  }
  
  Adapted from official impl at https://github.com/facebookresearch/LeViT, original copyright bellow.
  
  This version combines both conv/linear models and fixes torchscript compatibility.
  
  Modifications by/coyright Copyright 2021 Ross Wightman
  """
  
  # Copyright (c) 2015-present, Facebook, Inc.
  # All rights reserved.
  
  # Modified from
  # https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
  # Copyright 2020 Ross Wightman, Apache-2.0 License
  import itertools
  from copy import deepcopy
  from functools import partial
  
  import torch
  import torch.nn as nn
  
  from timm.data import IMAGENET_DEFAULT_STD, IMAGENET_DEFAULT_MEAN
  from .helpers import build_model_with_cfg, overlay_external_default_cfg
  from .layers import to_ntuple
  from .vision_transformer import trunc_normal_
  from .registry import register_model
  
  
  def _cfg(url='', **kwargs):
      return {
          'url': url,
          'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
          'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,
          'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
          'first_conv': 'patch_embed.0.c', 'classifier': ('head.l', 'head_dist.l'),
          **kwargs
      }
  
  
  default_cfgs = dict(
      levit_128s=_cfg(
          url='https://dl.fbaipublicfiles.com/LeViT/LeViT-128S-96703c44.pth'
      ),
      levit_128=_cfg(
          url='https://dl.fbaipublicfiles.com/LeViT/LeViT-128-b88c2750.pth'
      ),
      levit_192=_cfg(
          url='https://dl.fbaipublicfiles.com/LeViT/LeViT-192-92712e41.pth'
      ),
      levit_256=_cfg(
          url='https://dl.fbaipublicfiles.com/LeViT/LeViT-256-13b5763e.pth'
      ),
      levit_384=_cfg(
          url='https://dl.fbaipublicfiles.com/LeViT/LeViT-384-9bdaf2e2.pth'
      ),
  )
  
  model_cfgs = dict(
      levit_128s=dict(
          embed_dim=(128, 256, 384), key_dim=16, num_heads=(4, 6, 8), depth=(2, 3, 4)),
      levit_128=dict(
          embed_dim=(128, 256, 384), key_dim=16, num_heads=(4, 8, 12), depth=(4, 4, 4)),
      levit_192=dict(
          embed_dim=(192, 288, 384), key_dim=32, num_heads=(3, 5, 6), depth=(4, 4, 4)),
      levit_256=dict(
          embed_dim=(256, 384, 512), key_dim=32, num_heads=(4, 6, 8), depth=(4, 4, 4)),
      levit_384=dict(
          embed_dim=(384, 512, 768), key_dim=32, num_heads=(6, 9, 12), depth=(4, 4, 4)),
  )
  
  __all__ = ['Levit']
  
  
  @register_model
  def levit_128s(pretrained=False, fuse=False,distillation=True, use_conv=False, **kwargs):
      return create_levit(
          'levit_128s', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_128(pretrained=False, fuse=False, distillation=True, use_conv=False, **kwargs):
      return create_levit(
          'levit_128', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_192(pretrained=False, fuse=False, distillation=True, use_conv=False, **kwargs):
      return create_levit(
          'levit_192', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_256(pretrained=False, fuse=False, distillation=True, use_conv=False, **kwargs):
      return create_levit(
          'levit_256', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_384(pretrained=False, fuse=False, distillation=True, use_conv=False, **kwargs):
      return create_levit(
          'levit_384', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_c_128s(pretrained=False, fuse=False, distillation=True, use_conv=True,**kwargs):
      return create_levit(
          'levit_128s', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_c_128(pretrained=False, fuse=False,distillation=True, use_conv=True,  **kwargs):
      return create_levit(
          'levit_128', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_c_192(pretrained=False, fuse=False, distillation=True, use_conv=True, **kwargs):
      return create_levit(
          'levit_192', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_c_256(pretrained=False, fuse=False, distillation=True, use_conv=True, **kwargs):
      return create_levit(
          'levit_256', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  @register_model
  def levit_c_384(pretrained=False, fuse=False, distillation=True, use_conv=True, **kwargs):
      return create_levit(
          'levit_384', pretrained=pretrained, fuse=fuse,  distillation=distillation, use_conv=use_conv, **kwargs)
  
  
  class ConvNorm(nn.Sequential):
      def __init__(
              self, a, b, ks=1, stride=1, pad=0, dilation=1, groups=1, bn_weight_init=1, resolution=-10000):
          super().__init__()
          self.add_module('c', nn.Conv2d(a, b, ks, stride, pad, dilation, groups, bias=False))
          bn = nn.BatchNorm2d(b)
          nn.init.constant_(bn.weight, bn_weight_init)
          nn.init.constant_(bn.bias, 0)
          self.add_module('bn', bn)
  
      @torch.no_grad()
      def fuse(self):
          c, bn = self._modules.values()
          w = bn.weight / (bn.running_var + bn.eps) ** 0.5
          w = c.weight * w[:, None, None, None]
          b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
          m = nn.Conv2d(
              w.size(1), w.size(0), w.shape[2:], stride=self.c.stride,
              padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups)
          m.weight.data.copy_(w)
          m.bias.data.copy_(b)
          return m
  
  
  class LinearNorm(nn.Sequential):
      def __init__(self, a, b, bn_weight_init=1, resolution=-100000):
          super().__init__()
          self.add_module('c', nn.Linear(a, b, bias=False))
          bn = nn.BatchNorm1d(b)
          nn.init.constant_(bn.weight, bn_weight_init)
          nn.init.constant_(bn.bias, 0)
          self.add_module('bn', bn)
  
      @torch.no_grad()
      def fuse(self):
          l, bn = self._modules.values()
          w = bn.weight / (bn.running_var + bn.eps) ** 0.5
          w = l.weight * w[:, None]
          b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
          m = nn.Linear(w.size(1), w.size(0))
          m.weight.data.copy_(w)
          m.bias.data.copy_(b)
          return m
  
      def forward(self, x):
          x = self.c(x)
          return self.bn(x.flatten(0, 1)).reshape_as(x)
  
  
  class NormLinear(nn.Sequential):
      def __init__(self, a, b, bias=True, std=0.02):
          super().__init__()
          self.add_module('bn', nn.BatchNorm1d(a))
          l = nn.Linear(a, b, bias=bias)
          trunc_normal_(l.weight, std=std)
          if bias:
              nn.init.constant_(l.bias, 0)
          self.add_module('l', l)
  
      @torch.no_grad()
      def fuse(self):
          bn, l = self._modules.values()
          w = bn.weight / (bn.running_var + bn.eps) ** 0.5
          b = bn.bias - self.bn.running_mean * self.bn.weight / (bn.running_var + bn.eps) ** 0.5
          w = l.weight * w[None, :]
          if l.bias is None:
              b = b @ self.l.weight.T
          else:
              b = (l.weight @ b[:, None]).view(-1) + self.l.bias
          m = nn.Linear(w.size(1), w.size(0))
          m.weight.data.copy_(w)
          m.bias.data.copy_(b)
          return m
  
  
  def stem_b16(in_chs, out_chs, activation, resolution=224):
      return nn.Sequential(
          ConvNorm(in_chs, out_chs // 8, 3, 2, 1, resolution=resolution),
          activation(),
          ConvNorm(out_chs // 8, out_chs // 4, 3, 2, 1, resolution=resolution // 2),
          activation(),
          ConvNorm(out_chs // 4, out_chs // 2, 3, 2, 1, resolution=resolution // 4),
          activation(),
          ConvNorm(out_chs // 2, out_chs, 3, 2, 1, resolution=resolution // 8))
  
  
  class Residual(nn.Module):
      def __init__(self, m, drop):
          super().__init__()
          self.m = m
          self.drop = drop
  
      def forward(self, x):
          if self.training and self.drop > 0:
              return x + self.m(x) * torch.rand(
                  x.size(0), 1, 1, device=x.device).ge_(self.drop).div(1 - self.drop).detach()
          else:
              return x + self.m(x)
  
  
  class Subsample(nn.Module):
      def __init__(self, stride, resolution):
          super().__init__()
          self.stride = stride
          self.resolution = resolution
  
      def forward(self, x):
          B, N, C = x.shape
          x = x.view(B, self.resolution, self.resolution, C)[:, ::self.stride, ::self.stride]
          return x.reshape(B, -1, C)
  
  
  class Attention(nn.Module):
      def __init__(
              self, dim, key_dim, num_heads=8, attn_ratio=4, act_layer=None, resolution=14, use_conv=False):
          super().__init__()
  
          self.num_heads = num_heads
          self.scale = key_dim ** -0.5
          self.key_dim = key_dim
          self.nh_kd = nh_kd = key_dim * num_heads
          self.d = int(attn_ratio * key_dim)
          self.dh = int(attn_ratio * key_dim) * num_heads
          self.attn_ratio = attn_ratio
          self.use_conv = use_conv
          ln_layer = ConvNorm if self.use_conv else LinearNorm
          h = self.dh + nh_kd * 2
          self.qkv = ln_layer(dim, h, resolution=resolution)
          self.proj = nn.Sequential(
              act_layer(),
              ln_layer(self.dh, dim, bn_weight_init=0, resolution=resolution))
  
          points = list(itertools.product(range(resolution), range(resolution)))
          N = len(points)
          attention_offsets = {}
          idxs = []
          for p1 in points:
              for p2 in points:
                  offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
                  if offset not in attention_offsets:
                      attention_offsets[offset] = len(attention_offsets)
                  idxs.append(attention_offsets[offset])
          self.attention_biases = nn.Parameter(torch.zeros(num_heads, len(attention_offsets)))
          self.register_buffer('attention_bias_idxs', torch.LongTensor(idxs).view(N, N))
          self.ab = None
  
      @torch.no_grad()
      def train(self, mode=True):
          super().train(mode)
          self.ab = None if mode else self.attention_biases[:, self.attention_bias_idxs]
  
      def forward(self, x):  # x (B,C,H,W)
          if self.use_conv:
              B, C, H, W = x.shape
              q, k, v = self.qkv(x).view(B, self.num_heads, -1, H * W).split([self.key_dim, self.key_dim, self.d], dim=2)
              ab = self.attention_biases[:, self.attention_bias_idxs] if self.ab is None else self.ab
              attn = (q.transpose(-2, -1) @ k) * self.scale + ab
              attn = attn.softmax(dim=-1)
              x = (v @ attn.transpose(-2, -1)).view(B, -1, H, W)
          else:
              B, N, C = x.shape
              qkv = self.qkv(x)
              q, k, v = qkv.view(B, N, self.num_heads, -1).split([self.key_dim, self.key_dim, self.d], dim=3)
              q = q.permute(0, 2, 1, 3)
              k = k.permute(0, 2, 1, 3)
              v = v.permute(0, 2, 1, 3)
              ab = self.attention_biases[:, self.attention_bias_idxs] if self.ab is None else self.ab
              attn = q @ k.transpose(-2, -1) * self.scale + ab
              attn = attn.softmax(dim=-1)
              x = (attn @ v).transpose(1, 2).reshape(B, N, self.dh)
          x = self.proj(x)
          return x
  
  
  class AttentionSubsample(nn.Module):
      def __init__(
              self, in_dim, out_dim, key_dim, num_heads=8, attn_ratio=2,
              act_layer=None, stride=2, resolution=14, resolution_=7, use_conv=False):
          super().__init__()
          self.num_heads = num_heads
          self.scale = key_dim ** -0.5
          self.key_dim = key_dim
          self.nh_kd = nh_kd = key_dim * num_heads
          self.d = int(attn_ratio * key_dim)
          self.dh = self.d * self.num_heads
          self.attn_ratio = attn_ratio
          self.resolution_ = resolution_
          self.resolution_2 = resolution_ ** 2
          self.use_conv = use_conv
          if self.use_conv:
              ln_layer = ConvNorm
              sub_layer = partial(nn.AvgPool2d, kernel_size=1, padding=0)
          else:
              ln_layer = LinearNorm
              sub_layer = partial(Subsample, resolution=resolution)
  
          h = self.dh + nh_kd
          self.kv = ln_layer(in_dim, h, resolution=resolution)
          self.q = nn.Sequential(
              sub_layer(stride=stride),
              ln_layer(in_dim, nh_kd, resolution=resolution_))
          self.proj = nn.Sequential(
              act_layer(),
              ln_layer(self.dh, out_dim, resolution=resolution_))
  
          self.stride = stride
          self.resolution = resolution
          points = list(itertools.product(range(resolution), range(resolution)))
          points_ = list(itertools.product(range(resolution_), range(resolution_)))
          N = len(points)
          N_ = len(points_)
          attention_offsets = {}
          idxs = []
          for p1 in points_:
              for p2 in points:
                  size = 1
                  offset = (
                      abs(p1[0] * stride - p2[0] + (size - 1) / 2),
                      abs(p1[1] * stride - p2[1] + (size - 1) / 2))
                  if offset not in attention_offsets:
                      attention_offsets[offset] = len(attention_offsets)
                  idxs.append(attention_offsets[offset])
          self.attention_biases = nn.Parameter(torch.zeros(num_heads, len(attention_offsets)))
          self.register_buffer('attention_bias_idxs', torch.LongTensor(idxs).view(N_, N))
          self.ab = None
  
      @torch.no_grad()
      def train(self, mode=True):
          super().train(mode)
          self.ab = None if mode else self.attention_biases[:, self.attention_bias_idxs]
  
      def forward(self, x):
          if self.use_conv:
              B, C, H, W = x.shape
              k, v = self.kv(x).view(B, self.num_heads, -1, H * W).split([self.key_dim, self.d], dim=2)
              q = self.q(x).view(B, self.num_heads, self.key_dim, self.resolution_2)
  
              ab = self.attention_biases[:, self.attention_bias_idxs] if self.ab is None else self.ab
              attn = (q.transpose(-2, -1) @ k) * self.scale + ab
              attn = attn.softmax(dim=-1)
  
              x = (v @ attn.transpose(-2, -1)).reshape(B, -1, self.resolution_, self.resolution_)
          else:
              B, N, C = x.shape
              k, v = self.kv(x).view(B, N, self.num_heads, -1).split([self.key_dim, self.d], dim=3)
              k = k.permute(0, 2, 1, 3)  # BHNC
              v = v.permute(0, 2, 1, 3)  # BHNC
              q = self.q(x).view(B, self.resolution_2, self.num_heads, self.key_dim).permute(0, 2, 1, 3)
  
              ab = self.attention_biases[:, self.attention_bias_idxs] if self.ab is None else self.ab
              attn = q @ k.transpose(-2, -1) * self.scale + ab
              attn = attn.softmax(dim=-1)
  
              x = (attn @ v).transpose(1, 2).reshape(B, -1, self.dh)
          x = self.proj(x)
          return x
  
  
  class Levit(nn.Module):
      """ Vision Transformer with support for patch or hybrid CNN input stage
      """
  
      def __init__(
              self,
              img_size=224,
              patch_size=16,
              in_chans=3,
              num_classes=1000,
              embed_dim=(192,),
              key_dim=64,
              depth=(12,),
              num_heads=(3,),
              attn_ratio=2,
              mlp_ratio=2,
              hybrid_backbone=None,
              down_ops=None,
              act_layer=nn.Hardswish,
              attn_act_layer=nn.Hardswish,
              distillation=True,
              use_conv=False,
              drop_path=0):
          super().__init__()
          if isinstance(img_size, tuple):
              # FIXME origin impl passes single img/res dim through whole hierarchy,
              # not sure this model will be used enough to spend time fixing it.
              assert img_size[0] == img_size[1]
              img_size = img_size[0]
          self.num_classes = num_classes
          self.num_features = embed_dim[-1]
          self.embed_dim = embed_dim
          N = len(embed_dim)
          assert len(depth) == len(num_heads) == N
          key_dim = to_ntuple(N)(key_dim)
          attn_ratio = to_ntuple(N)(attn_ratio)
          mlp_ratio = to_ntuple(N)(mlp_ratio)
          down_ops = down_ops or (
              # ('Subsample',key_dim, num_heads, attn_ratio, mlp_ratio, stride)
              ('Subsample', key_dim[0], embed_dim[0] // key_dim[0], 4, 2, 2),
              ('Subsample', key_dim[0], embed_dim[1] // key_dim[1], 4, 2, 2),
              ('',)
          )
          self.distillation = distillation
          self.use_conv = use_conv
          ln_layer = ConvNorm if self.use_conv else LinearNorm
  
          self.patch_embed = hybrid_backbone or stem_b16(in_chans, embed_dim[0], activation=act_layer)
  
          self.blocks = []
          resolution = img_size // patch_size
          for i, (ed, kd, dpth, nh, ar, mr, do) in enumerate(
                  zip(embed_dim, key_dim, depth, num_heads, attn_ratio, mlp_ratio, down_ops)):
              for _ in range(dpth):
                  self.blocks.append(
                      Residual(
                          Attention(
                              ed, kd, nh, attn_ratio=ar, act_layer=attn_act_layer,
                              resolution=resolution, use_conv=use_conv),
                          drop_path))
                  if mr > 0:
                      h = int(ed * mr)
                      self.blocks.append(
                          Residual(nn.Sequential(
                              ln_layer(ed, h, resolution=resolution),
                              act_layer(),
                              ln_layer(h, ed, bn_weight_init=0, resolution=resolution),
                          ), drop_path))
              if do[0] == 'Subsample':
                  # ('Subsample',key_dim, num_heads, attn_ratio, mlp_ratio, stride)
                  resolution_ = (resolution - 1) // do[5] + 1
                  self.blocks.append(
                      AttentionSubsample(
                          *embed_dim[i:i + 2], key_dim=do[1], num_heads=do[2],
                          attn_ratio=do[3], act_layer=attn_act_layer, stride=do[5],
                          resolution=resolution, resolution_=resolution_, use_conv=use_conv))
                  resolution = resolution_
                  if do[4] > 0:  # mlp_ratio
                      h = int(embed_dim[i + 1] * do[4])
                      self.blocks.append(
                          Residual(nn.Sequential(
                              ln_layer(embed_dim[i + 1], h, resolution=resolution),
                              act_layer(),
                              ln_layer(h, embed_dim[i + 1], bn_weight_init=0, resolution=resolution),
                          ), drop_path))
          self.blocks = nn.Sequential(*self.blocks)
  
          # Classifier head
          self.head = NormLinear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity()
          if distillation:
              self.head_dist = NormLinear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity()
          else:
              self.head_dist = None
  
      @torch.jit.ignore
      def no_weight_decay(self):
          return {x for x in self.state_dict().keys() if 'attention_biases' in x}
  
      def forward(self, x):
          x = self.patch_embed(x)
          if not self.use_conv:
              x = x.flatten(2).transpose(1, 2)
          x = self.blocks(x)
          x = x.mean((-2, -1)) if self.use_conv else x.mean(1)
          if self.head_dist is not None:
              x, x_dist = self.head(x), self.head_dist(x)
              if self.training and not torch.jit.is_scripting():
                  return x, x_dist
              else:
                  # during inference, return the average of both classifier predictions
                  return (x + x_dist) / 2
          else:
              x = self.head(x)
          return x
  
  
  def checkpoint_filter_fn(state_dict, model):
      if 'model' in state_dict:
          # For deit models
          state_dict = state_dict['model']
      D = model.state_dict()
      for k in state_dict.keys():
          if D[k].ndim == 4 and state_dict[k].ndim == 2:
              state_dict[k] = state_dict[k][:, :, None, None]
      return state_dict
  
  
  def create_levit(variant, pretrained=False, default_cfg=None, fuse=False, **kwargs):
      if kwargs.get('features_only', None):
          raise RuntimeError('features_only not implemented for Vision Transformer models.')
  
      model_cfg = dict(**model_cfgs[variant], **kwargs)
      model = build_model_with_cfg(
          Levit, variant, pretrained,
          default_cfg=default_cfgs[variant],
          pretrained_filter_fn=checkpoint_filter_fn,
          **model_cfg)
      #if fuse:
      #    utils.replace_batchnorm(model)
      return model
  
  