def __init__(
        self,
        data_info: dict,
        embed_size: int = 300,
        num_layers: int = 5,
        dropout: float = 0.5,
        virtual_node: bool = False,
    ):
        """Graph Isomorphism Network (GIN) variant introduced in baselines
        for OGB graph property prediction datasets

        Parameters
        ----------
        data_info : dict
            The information about the input dataset.
        embed_size : int
            Embedding size.
        num_layers : int
            Number of layers.
        dropout : float
            Dropout rate.
        virtual_node : bool
            Whether to use virtual node.
        """
        super(OGBGGIN, self).__init__()
        self.data_info = data_info
        self.embed_size = embed_size
        self.num_layers = num_layers
        self.virtual_node = virtual_node

        if data_info["name"] in ["ogbg-molhiv", "ogbg-molpcba"]:
            self.node_encoder = AtomEncoder(embed_size)
            self.edge_encoders = nn.ModuleList(
                [BondEncoder(embed_size) for _ in range(num_layers)]
            )
        else:
            # Handle other datasets
            self.node_encoder = nn.Linear(
                data_info["node_feat_size"], embed_size
            )
            self.edge_encoders = nn.ModuleList(
                [
                    nn.Linear(data_info["edge_feat_size"], embed_size)
                    for _ in range(num_layers)
                ]
            )

        self.conv_layers = nn.ModuleList(
            [GINEConv(MLP(embed_size)) for _ in range(num_layers)]
        )

        self.dropout = nn.Dropout(dropout)
        self.pool = AvgPooling()
        self.pred = nn.Linear(embed_size, data_info["out_size"])

        if virtual_node:
            self.virtual_emb = nn.Embedding(1, embed_size)
            nn.init.constant_(self.virtual_emb.weight.data, 0)
            self.mlp_virtual = nn.ModuleList()
            for _ in range(num_layers - 1):
                self.mlp_virtual.append(MLP(embed_size))
            self.virtual_pool = SumPooling()
Graph Isomorphism Network (GIN) variant introduced in baselines
        for OGB graph property prediction datasets

        Parameters
        ----------
        data_info : dict
            The information about the input dataset.
        embed_size : int
            Embedding size.
        num_layers : int
            Number of layers.
        dropout : float
            Dropout rate.
        virtual_node : bool
            Whether to use virtual node.
        
dglgo/dglgo/model/graph_encoder/gin_ogbg.py