def __init__(
        self,
        in_channels,
        out_channels,
        fast_gelu=True,
        dtype="float16",
    ):
        """Initialize dual gemm module, create a tensor for weights"""
        super().__init__()
        self.w1 = Parameter(shape=[out_channels, in_channels], dtype=dtype)
        self.w2 = Parameter(shape=[out_channels, in_channels], dtype=dtype)
        if fast_gelu:
            self.op = ops.dual_gemm_rcr_fast_gelu()
        else:
            self.op = ops.dual_gemm_rcr_silu()
Initialize dual gemm module, create a tensor for weights
python/aitemplate/frontend/nn/dual_gemm.py