def __init__(
        self,
        sequence_pool: Module = None,
        dropout: Module = None,
        proj: Module = None,
        activation: Module = None,
    ) -> None:
        """
        Args:
            sequence_pool (torch.nn.modules): pooling module.
            dropout(torch.nn.modules): dropout module.
            proj (torch.nn.modules): project module.
            activation (torch.nn.modules): activation module.
        """
        super().__init__()
        self.sequence_pool = sequence_pool
        self.dropout = dropout
        self.proj = proj
        self.activation = activation

        Args:
            sequence_pool (torch.nn.modules): pooling module.
            dropout(torch.nn.modules): dropout module.
            proj (torch.nn.modules): project module.
            activation (torch.nn.modules): activation module.
        
python/aitemplate/frontend/nn/head.py