def __init__(
        self, input_shape: ShapeSpec, *, conv_dim: int, fc_dims: List[int], output_shape: Tuple[int]
    ):
        """
        Args:
            conv_dim: the output dimension of the conv layers
            fc_dims: a list of N>0 integers representing the output dimensions of N FC layers
            output_shape: shape of the output mask prediction
        """
        super().__init__()

        # fmt: off
        input_channels    = input_shape.channels
        input_h           = input_shape.height
        input_w           = input_shape.width
        self.output_shape = output_shape
        # fmt: on

        self.conv_layers = []
        if input_channels > conv_dim:
            self.reduce_channel_dim_conv = Conv2d(
                input_channels,
                conv_dim,
                kernel_size=1,
                stride=1,
                padding=0,
                bias=True,
                activation=F.relu,
            )
            self.conv_layers.append(self.reduce_channel_dim_conv)

        self.reduce_spatial_dim_conv = Conv2d(
            conv_dim, conv_dim, kernel_size=2, stride=2, padding=0, bias=True, activation=F.relu
        )
        self.conv_layers.append(self.reduce_spatial_dim_conv)

        input_dim = conv_dim * input_h * input_w
        input_dim //= 4

        self.fcs = []
        for k, fc_dim in enumerate(fc_dims):
            fc = nn.Linear(input_dim, fc_dim)
            self.add_module("fc{}".format(k + 1), fc)
            self.fcs.append(fc)
            input_dim = fc_dim

        output_dim = int(np.prod(self.output_shape))

        self.prediction = nn.Linear(fc_dims[-1], output_dim)
        # use normal distribution initialization for mask prediction layer
        nn.init.normal_(self.prediction.weight, std=0.001)
        nn.init.constant_(self.prediction.bias, 0)

        for layer in self.conv_layers:
            weight_init.c2_msra_fill(layer)
        for layer in self.fcs:
            weight_init.c2_xavier_fill(layer)

        Args:
            conv_dim: the output dimension of the conv layers
            fc_dims: a list of N>0 integers representing the output dimensions of N FC layers
            output_shape: shape of the output mask prediction
        
EVA-02/det/projects/PointRend/point_rend/mask_head.py