import torch
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice
import logging
import torchvision
import warnings
from collections import OrderedDict
from torch.utils import model_zoo
from torch.nn import functional as F
import torch.nn as nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused__to_copy__unsafe_index_add_arange_clamp_floor_mul_rsub_sub_0(
    in_out_ptr1, in_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 3072
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 16 % 16
    x0 = xindex % 16
    x2 = xindex // 256
    x3 = xindex
    tmp0 = x1
    tmp1 = tmp0.to(tl.float32)
    tmp2 = 0.5
    tmp3 = tmp1 + tmp2
    tmp4 = 0.25
    tmp5 = tmp3 * tmp4
    tmp6 = tmp5 - tmp2
    tmp7 = libdevice.floor(tmp6)
    tmp8 = tmp7.to(tl.int32)
    tmp9 = tl.full([1], 1, tl.int64)
    tmp10 = tmp8 - tmp9
    tmp11 = tl.full([1], 0, tl.int64)
    tmp12 = triton_helpers.maximum(tmp10, tmp11)
    tmp13 = tl.full([1], 3, tl.int64)
    tmp14 = triton_helpers.minimum(tmp12, tmp13)
    tmp15 = x0
    tmp16 = tmp15.to(tl.float32)
    tmp17 = tmp16 + tmp2
    tmp18 = tmp17 * tmp4
    tmp19 = tmp18 - tmp2
    tmp20 = libdevice.floor(tmp19)
    tmp21 = tmp20.to(tl.int32)
    tmp22 = tmp21 - tmp9
    tmp23 = triton_helpers.maximum(tmp22, tmp11)
    tmp24 = triton_helpers.minimum(tmp23, tmp13)
    tmp25 = tl.load(in_ptr0 + (tmp24 + 4 * tmp14 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp26 = tmp19 - tmp20
    tmp27 = 0.0
    tmp28 = triton_helpers.maximum(tmp26, tmp27)
    tmp29 = 1.0
    tmp30 = triton_helpers.minimum(tmp28, tmp29)
    tmp31 = tmp30 + tmp29
    tmp32 = -0.75
    tmp33 = tmp31 * tmp32
    tmp34 = -3.75
    tmp35 = tmp33 - tmp34
    tmp36 = tmp35 * tmp31
    tmp37 = -6.0
    tmp38 = tmp36 + tmp37
    tmp39 = tmp38 * tmp31
    tmp40 = -3.0
    tmp41 = tmp39 - tmp40
    tmp42 = tmp25 * tmp41
    tmp43 = triton_helpers.maximum(tmp21, tmp11)
    tmp44 = triton_helpers.minimum(tmp43, tmp13)
    tmp45 = tl.load(in_ptr0 + (tmp44 + 4 * tmp14 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp46 = 1.25
    tmp47 = tmp30 * tmp46
    tmp48 = 2.25
    tmp49 = tmp47 - tmp48
    tmp50 = tmp49 * tmp30
    tmp51 = tmp50 * tmp30
    tmp52 = tmp51 + tmp29
    tmp53 = tmp45 * tmp52
    tmp54 = tmp21 + tmp9
    tmp55 = triton_helpers.maximum(tmp54, tmp11)
    tmp56 = triton_helpers.minimum(tmp55, tmp13)
    tmp57 = tl.load(in_ptr0 + (tmp56 + 4 * tmp14 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp58 = tmp29 - tmp30
    tmp59 = tmp58 * tmp46
    tmp60 = tmp59 - tmp48
    tmp61 = tmp60 * tmp58
    tmp62 = tmp61 * tmp58
    tmp63 = tmp62 + tmp29
    tmp64 = tmp57 * tmp63
    tmp65 = triton_helpers.maximum(tmp8, tmp11)
    tmp66 = triton_helpers.minimum(tmp65, tmp13)
    tmp67 = tl.load(in_ptr0 + (tmp24 + 4 * tmp66 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp68 = tmp67 * tmp41
    tmp69 = tl.full([1], 2, tl.int64)
    tmp70 = tmp21 + tmp69
    tmp71 = triton_helpers.maximum(tmp70, tmp11)
    tmp72 = triton_helpers.minimum(tmp71, tmp13)
    tmp73 = tl.load(in_ptr0 + (tmp72 + 4 * tmp14 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp74 = 2.0
    tmp75 = tmp74 - tmp30
    tmp76 = tmp75 * tmp32
    tmp77 = tmp76 - tmp34
    tmp78 = tmp77 * tmp75
    tmp79 = tmp78 + tmp37
    tmp80 = tmp79 * tmp75
    tmp81 = tmp80 - tmp40
    tmp82 = tmp73 * tmp81
    tmp83 = tl.load(in_ptr0 + (tmp44 + 4 * tmp66 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp84 = tmp83 * tmp52
    tmp85 = tl.load(in_ptr0 + (tmp56 + 4 * tmp66 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp86 = tmp85 * tmp63
    tmp87 = tmp8 + tmp9
    tmp88 = triton_helpers.maximum(tmp87, tmp11)
    tmp89 = triton_helpers.minimum(tmp88, tmp13)
    tmp90 = tl.load(in_ptr0 + (tmp24 + 4 * tmp89 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp91 = tmp90 * tmp41
    tmp92 = tl.load(in_ptr0 + (tmp72 + 4 * tmp66 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp93 = tmp92 * tmp81
    tmp94 = tl.load(in_ptr0 + (tmp44 + 4 * tmp89 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp95 = tmp94 * tmp52
    tmp96 = tl.load(in_ptr0 + (tmp56 + 4 * tmp89 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp97 = tmp96 * tmp63
    tmp98 = tmp8 + tmp69
    tmp99 = triton_helpers.maximum(tmp98, tmp11)
    tmp100 = triton_helpers.minimum(tmp99, tmp13)
    tmp101 = tl.load(in_ptr0 + (tmp24 + 4 * tmp100 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp102 = tmp101 * tmp41
    tmp103 = tl.load(in_ptr0 + (tmp72 + 4 * tmp89 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp104 = tmp103 * tmp81
    tmp105 = tl.load(in_ptr0 + (tmp44 + 4 * tmp100 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp106 = tmp105 * tmp52
    tmp107 = tl.load(in_ptr0 + (tmp56 + 4 * tmp100 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp108 = tmp107 * tmp63
    tmp109 = tl.load(in_ptr0 + (tmp72 + 4 * tmp100 + 16 * x2), xmask,
        eviction_policy='evict_last')
    tmp110 = tmp109 * tmp81
    tmp111 = tmp42 + tmp53
    tmp112 = tmp111 + tmp64
    tmp113 = tmp112 + tmp82
    tmp114 = tmp6 - tmp7
    tmp115 = triton_helpers.maximum(tmp114, tmp27)
    tmp116 = triton_helpers.minimum(tmp115, tmp29)
    tmp117 = tmp116 + tmp29
    tmp118 = tmp117 * tmp32
    tmp119 = tmp118 - tmp34
    tmp120 = tmp119 * tmp117
    tmp121 = tmp120 + tmp37
    tmp122 = tmp121 * tmp117
    tmp123 = tmp122 - tmp40
    tmp124 = tmp113 * tmp123
    tmp125 = tmp68 + tmp84
    tmp126 = tmp125 + tmp86
    tmp127 = tmp126 + tmp93
    tmp128 = tmp116 * tmp46
    tmp129 = tmp128 - tmp48
    tmp130 = tmp129 * tmp116
    tmp131 = tmp130 * tmp116
    tmp132 = tmp131 + tmp29
    tmp133 = tmp127 * tmp132
    tmp134 = tmp124 + tmp133
    tmp135 = tmp91 + tmp95
    tmp136 = tmp135 + tmp97
    tmp137 = tmp136 + tmp104
    tmp138 = tmp29 - tmp116
    tmp139 = tmp138 * tmp46
    tmp140 = tmp139 - tmp48
    tmp141 = tmp140 * tmp138
    tmp142 = tmp141 * tmp138
    tmp143 = tmp142 + tmp29
    tmp144 = tmp137 * tmp143
    tmp145 = tmp134 + tmp144
    tmp146 = tmp102 + tmp106
    tmp147 = tmp146 + tmp108
    tmp148 = tmp147 + tmp110
    tmp149 = tmp74 - tmp116
    tmp150 = tmp149 * tmp32
    tmp151 = tmp150 - tmp34
    tmp152 = tmp151 * tmp149
    tmp153 = tmp152 + tmp37
    tmp154 = tmp153 * tmp149
    tmp155 = tmp154 - tmp40
    tmp156 = tmp148 * tmp155
    tmp157 = tmp145 + tmp156
    tl.store(in_out_ptr1 + x3, tmp157, xmask)


@triton.jit
def triton_poi_fused_convolution_relu_1(in_out_ptr0, in_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x3 = xindex
    x1 = xindex // 256 % 64
    tmp0 = tl.load(in_out_ptr0 + x3, None)
    tmp1 = tl.load(in_ptr0 + x1, None, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tl.store(in_out_ptr0 + x3, tmp4, None)


@triton.jit
def triton_poi_fused_convolution_relu_2(in_out_ptr0, in_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x3 = xindex
    x1 = xindex // 256 % 32
    tmp0 = tl.load(in_out_ptr0 + x3, None)
    tmp1 = tl.load(in_ptr0 + x1, None, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tl.store(in_out_ptr0 + x3, tmp4, None)


@triton.jit
def triton_poi_fused_convolution_3(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 3072
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x1 = xindex // 256 % 3
    tmp0 = tl.load(in_out_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(in_out_ptr0 + x3, tmp2, xmask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7) = args
    args.clear()
    assert_size_stride(primals_1, (4, 3, 4, 4), (48, 16, 4, 1))
    assert_size_stride(primals_2, (64, 3, 9, 9), (243, 81, 9, 1))
    assert_size_stride(primals_3, (64,), (1,))
    assert_size_stride(primals_4, (32, 64, 1, 1), (64, 1, 1, 1))
    assert_size_stride(primals_5, (32,), (1,))
    assert_size_stride(primals_6, (3, 32, 5, 5), (800, 25, 5, 1))
    assert_size_stride(primals_7, (3,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf10 = empty_strided_cuda((4, 3, 16, 16), (768, 256, 16, 1), torch
            .float32)
        buf18 = buf10
        del buf10
        buf20 = buf18
        del buf18
        get_raw_stream(0)
        triton_poi_fused__to_copy__unsafe_index_add_arange_clamp_floor_mul_rsub_sub_0[
            grid(3072)](buf20, primals_1, 3072, XBLOCK=128, num_warps=4,
            num_stages=1)
        del primals_1
        buf21 = extern_kernels.convolution(buf20, primals_2, stride=(1, 1),
            padding=(4, 4), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf21, (4, 64, 16, 16), (16384, 256, 16, 1))
        buf22 = buf21
        del buf21
        triton_poi_fused_convolution_relu_1[grid(65536)](buf22, primals_3, 
            65536, XBLOCK=512, num_warps=4, num_stages=1)
        del primals_3
        buf23 = extern_kernels.convolution(buf22, primals_4, stride=(1, 1),
            padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf23, (4, 32, 16, 16), (8192, 256, 16, 1))
        buf24 = buf23
        del buf23
        triton_poi_fused_convolution_relu_2[grid(32768)](buf24, primals_5, 
            32768, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_5
        buf25 = extern_kernels.convolution(buf24, primals_6, stride=(1, 1),
            padding=(2, 2), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf25, (4, 3, 16, 16), (768, 256, 16, 1))
        buf26 = buf25
        del buf25
        triton_poi_fused_convolution_3[grid(3072)](buf26, primals_7, 3072,
            XBLOCK=256, num_warps=4, num_stages=1)
        del primals_7
    return buf26, primals_2, primals_4, primals_6, buf20, buf22, buf24


def get_root_logger(log_file=None, log_level=logging.INFO):
    """Get the root logger.

    The logger will be initialized if it has not been initialized. By default a
    StreamHandler will be added. If `log_file` is specified, a FileHandler will
    also be added. The name of the root logger is the top-level package name,
    e.g., "mmedit".

    Args:
        log_file (str | None): The log filename. If specified, a FileHandler
            will be added to the root logger.
        log_level (int): The root logger level. Note that only the process of
            rank 0 is affected, while other processes will set the level to
            "Error" and be silent most of the time.

    Returns:
        logging.Logger: The root logger.
    """
    logger = get_logger(__name__.split('.')[0], log_file, log_level)
    return logger


def _get_mmcv_home():
    mmcv_home = os.path.expanduser(os.getenv(ENV_MMCV_HOME, os.path.join(os
        .getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'mmcv')))
    mkdir_or_exist(mmcv_home)
    return mmcv_home


def _process_mmcls_checkpoint(checkpoint):
    state_dict = checkpoint['state_dict']
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        if k.startswith('backbone.'):
            new_state_dict[k[9:]] = v
    new_checkpoint = dict(state_dict=new_state_dict)
    return new_checkpoint


def get_deprecated_model_names():
    deprecate_json_path = osp.join(mmcv.__path__[0],
        'model_zoo/deprecated.json')
    deprecate_urls = load_file(deprecate_json_path)
    assert isinstance(deprecate_urls, dict)
    return deprecate_urls


def get_external_models():
    mmcv_home = _get_mmcv_home()
    default_json_path = osp.join(mmcv.__path__[0], 'model_zoo/open_mmlab.json')
    default_urls = load_file(default_json_path)
    assert isinstance(default_urls, dict)
    external_json_path = osp.join(mmcv_home, 'open_mmlab.json')
    if osp.exists(external_json_path):
        external_urls = load_file(external_json_path)
        assert isinstance(external_urls, dict)
        default_urls.update(external_urls)
    return default_urls


def get_mmcls_models():
    mmcls_json_path = osp.join(mmcv.__path__[0], 'model_zoo/mmcls.json')
    mmcls_urls = load_file(mmcls_json_path)
    return mmcls_urls


def get_torchvision_models():
    model_urls = dict()
    for _, name, ispkg in pkgutil.walk_packages(torchvision.models.__path__):
        if ispkg:
            continue
        _zoo = import_module(f'torchvision.models.{name}')
        if hasattr(_zoo, 'model_urls'):
            _urls = getattr(_zoo, 'model_urls')
            model_urls.update(_urls)
    return model_urls


def load_fileclient_dist(filename, backend, map_location):
    """In distributed setting, this function only download checkpoint at local
    rank 0."""
    rank, world_size = get_dist_info()
    rank = int(os.environ.get('LOCAL_RANK', rank))
    allowed_backends = ['ceph']
    if backend not in allowed_backends:
        raise ValueError(f'Load from Backend {backend} is not supported.')
    if rank == 0:
        fileclient = FileClient(backend=backend)
        buffer = io.BytesIO(fileclient.get(filename))
        checkpoint = torch.load(buffer, map_location=map_location)
    if world_size > 1:
        torch.distributed.barrier()
        if rank > 0:
            fileclient = FileClient(backend=backend)
            buffer = io.BytesIO(fileclient.get(filename))
            checkpoint = torch.load(buffer, map_location=map_location)
    return checkpoint


def load_url_dist(url, model_dir=None):
    """In distributed setting, this function only download checkpoint at local
    rank 0."""
    rank, world_size = get_dist_info()
    rank = int(os.environ.get('LOCAL_RANK', rank))
    if rank == 0:
        checkpoint = model_zoo.load_url(url, model_dir=model_dir)
    if world_size > 1:
        torch.distributed.barrier()
        if rank > 0:
            checkpoint = model_zoo.load_url(url, model_dir=model_dir)
    return checkpoint


def _load_checkpoint(filename, map_location=None):
    """Load checkpoint from somewhere (modelzoo, file, url).

    Args:
        filename (str): Accept local filepath, URL, ``torchvision://xxx``,
            ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for
            details.
        map_location (str | None): Same as :func:`torch.load`. Default: None.

    Returns:
        dict | OrderedDict: The loaded checkpoint. It can be either an
            OrderedDict storing model weights or a dict containing other
            information, which depends on the checkpoint.
    """
    if filename.startswith('modelzoo://'):
        warnings.warn(
            'The URL scheme of "modelzoo://" is deprecated, please use "torchvision://" instead'
            )
        model_urls = get_torchvision_models()
        model_name = filename[11:]
        checkpoint = load_url_dist(model_urls[model_name])
    elif filename.startswith('torchvision://'):
        model_urls = get_torchvision_models()
        model_name = filename[14:]
        checkpoint = load_url_dist(model_urls[model_name])
    elif filename.startswith('open-mmlab://'):
        model_urls = get_external_models()
        model_name = filename[13:]
        deprecated_urls = get_deprecated_model_names()
        if model_name in deprecated_urls:
            warnings.warn(
                f'open-mmlab://{model_name} is deprecated in favor of open-mmlab://{deprecated_urls[model_name]}'
                )
            model_name = deprecated_urls[model_name]
        model_url = model_urls[model_name]
        if model_url.startswith(('http://', 'https://')):
            checkpoint = load_url_dist(model_url)
        else:
            filename = osp.join(_get_mmcv_home(), model_url)
            if not osp.isfile(filename):
                raise IOError(f'{filename} is not a checkpoint file')
            checkpoint = torch.load(filename, map_location=map_location)
    elif filename.startswith('mmcls://'):
        model_urls = get_mmcls_models()
        model_name = filename[8:]
        checkpoint = load_url_dist(model_urls[model_name])
        checkpoint = _process_mmcls_checkpoint(checkpoint)
    elif filename.startswith(('http://', 'https://')):
        checkpoint = load_url_dist(filename)
    elif filename.startswith('pavi://'):
        model_path = filename[7:]
        checkpoint = load_pavimodel_dist(model_path, map_location=map_location)
    elif filename.startswith('s3://'):
        checkpoint = load_fileclient_dist(filename, backend='ceph',
            map_location=map_location)
    else:
        if not osp.isfile(filename):
            raise IOError(f'{filename} is not a checkpoint file')
        checkpoint = torch.load(filename, map_location=map_location)
    return checkpoint


def load_state_dict(module, state_dict, strict=False, logger=None):
    """Load state_dict to a module.

    This method is modified from :meth:`torch.nn.Module.load_state_dict`.
    Default value for ``strict`` is set to ``False`` and the message for
    param mismatch will be shown even if strict is False.

    Args:
        module (Module): Module that receives the state_dict.
        state_dict (OrderedDict): Weights.
        strict (bool): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``False``.
        logger (:obj:`logging.Logger`, optional): Logger to log the error
            message. If not specified, print function will be used.
    """
    unexpected_keys = []
    all_missing_keys = []
    err_msg = []
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        state_dict._metadata = metadata

    def load(module, prefix=''):
        if is_module_wrapper(module):
            module = module.module
        local_metadata = {} if metadata is None else metadata.get(prefix[:-
            1], {})
        module._load_from_state_dict(state_dict, prefix, local_metadata, 
            True, all_missing_keys, unexpected_keys, err_msg)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')
    load(module)
    load = None
    missing_keys = [key for key in all_missing_keys if 
        'num_batches_tracked' not in key]
    if unexpected_keys:
        err_msg.append(
            f"unexpected key in source state_dict: {', '.join(unexpected_keys)}\n"
            )
    if missing_keys:
        err_msg.append(
            f"missing keys in source state_dict: {', '.join(missing_keys)}\n")
    rank, _ = get_dist_info()
    if len(err_msg) > 0 and rank == 0:
        err_msg.insert(0,
            'The model and loaded state dict do not match exactly\n')
        err_msg = '\n'.join(err_msg)
        if strict:
            raise RuntimeError(err_msg)
        elif logger is not None:
            logger.warning(err_msg)
        else:
            None


def load_checkpoint(model, filename, map_location='cpu', strict=False,
    logger=None):
    """Load checkpoint from a file or URI.

    Args:
        model (Module): Module to load checkpoint.
        filename (str): Accept local filepath, URL, ``torchvision://xxx``,
            ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for
            details.
        map_location (str): Same as :func:`torch.load`.
        strict (bool): Whether to allow different params for the model and
            checkpoint.
        logger (:mod:`logging.Logger` or None): The logger for error message.

    Returns:
        dict or OrderedDict: The loaded checkpoint.
    """
    checkpoint = _load_checkpoint(filename, map_location)
    if not isinstance(checkpoint, dict):
        raise RuntimeError(f'No state_dict found in checkpoint file {filename}'
            )
    if 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
    elif 'model' in checkpoint:
        state_dict = checkpoint['model']
    else:
        state_dict = checkpoint
    if list(state_dict.keys())[0].startswith('module.'):
        state_dict = {k[7:]: v for k, v in state_dict.items()}
    if state_dict.get('absolute_pos_embed') is not None:
        absolute_pos_embed = state_dict['absolute_pos_embed']
        N1, L, C1 = absolute_pos_embed.size()
        N2, C2, H, W = model.absolute_pos_embed.size()
        if N1 != N2 or C1 != C2 or L != H * W:
            logger.warning('Error in loading absolute_pos_embed, pass')
        else:
            state_dict['absolute_pos_embed'] = absolute_pos_embed.view(N2,
                H, W, C2).permute(0, 3, 1, 2)
    relative_position_bias_table_keys = [k for k in state_dict.keys() if 
        'relative_position_bias_table' in k]
    for table_key in relative_position_bias_table_keys:
        table_pretrained = state_dict[table_key]
        table_current = model.state_dict()[table_key]
        L1, nH1 = table_pretrained.size()
        L2, nH2 = table_current.size()
        if nH1 != nH2:
            logger.warning(f'Error in loading {table_key}, pass')
        elif L1 != L2:
            S1 = int(L1 ** 0.5)
            S2 = int(L2 ** 0.5)
            table_pretrained_resized = F.interpolate(table_pretrained.
                permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2), mode=
                'bicubic')
            state_dict[table_key] = table_pretrained_resized.view(nH2, L2
                ).permute(1, 0)
    load_state_dict(model, state_dict, strict, logger)
    return checkpoint


class SRCNNNew(nn.Module):
    """SRCNN network structure for image super resolution.

    SRCNN has three conv layers. For each layer, we can define the
    `in_channels`, `out_channels` and `kernel_size`.
    The input image will first be upsampled with a bicubic upsampler, and then
    super-resolved in the HR spatial size.

    Paper: Learning a Deep Convolutional Network for Image Super-Resolution.

    Args:
        channels (tuple[int]): A tuple of channel numbers for each layer
            including channels of input and output . Default: (3, 64, 32, 3).
        kernel_sizes (tuple[int]): A tuple of kernel sizes for each conv layer.
            Default: (9, 1, 5).
        upscale_factor (int): Upsampling factor. Default: 4.
    """

    def __init__(self, channels=(3, 64, 32, 3), kernel_sizes=(9, 1, 5),
        upscale_factor=4):
        super().__init__()
        assert len(channels
            ) == 4, f'The length of channel tuple should be 4, but got {len(channels)}'
        assert len(kernel_sizes
            ) == 3, f'The length of kernel tuple should be 3, but got {len(kernel_sizes)}'
        self.upscale_factor = upscale_factor
        self.img_upsampler = nn.Upsample(scale_factor=self.upscale_factor,
            mode='bicubic', align_corners=False)
        self.conv1 = nn.Conv2d(channels[0], channels[1], kernel_size=
            kernel_sizes[0], padding=kernel_sizes[0] // 2)
        self.conv2 = nn.Conv2d(channels[1], channels[2], kernel_size=
            kernel_sizes[1], padding=kernel_sizes[1] // 2)
        self.conv3 = nn.Conv2d(channels[2], channels[3], kernel_size=
            kernel_sizes[2], padding=kernel_sizes[2] // 2)
        self.relu = nn.ReLU()

    def init_weights(self, pretrained=None, strict=True):
        """Init weights for models.

        Args:
            pretrained (str, optional): Path for pretrained weights. If given
                None, pretrained weights will not be loaded. Defaults to None.
            strict (boo, optional): Whether strictly load the pretrained model.
                Defaults to True.
        """
        if isinstance(pretrained, str):
            logger = get_root_logger()
            load_checkpoint(self, pretrained, strict=strict, logger=logger)
        elif pretrained is None:
            pass
        else:
            raise TypeError(
                f'"pretrained" must be a str or None. But received {type(pretrained)}.'
                )

    def forward(self, input_0):
        primals_2 = self.conv1.weight
        primals_3 = self.conv1.bias
        primals_4 = self.conv2.weight
        primals_5 = self.conv2.bias
        primals_6 = self.conv3.weight
        primals_7 = self.conv3.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7])
        return output[0]


Measured latency_ms: 0.036636