# Copyright (c) Megvii, Inc. and its affiliates.

import megengine as mge
import megengine.functional as F
import megengine.module as M

from .resnet import BasicBlock


class STN(M.Module):
    """spatial transformer networks from
    `"Spatial Transformer Networks" <https://arxiv.org/pdf/1506.02025.pdf>`_
    some detailed implements are highly simplified while good performance maintained
    """

    def __init__(self, input_size=112):
        assert input_size == 112, f"expected input_size == 112, got {input_size}"
        super().__init__()
        self.input_size = input_size
        self.stem = M.Sequential(
            
M.Conv2d(3, 8, kernel_size=3, stride=2, padding=1, bias=False)