from .gcn import GCN
from .gat import GAT
from .gin import GIN
from .gine import GINE
from .pna import PNA
from .diff_pool import DiffPool
from gcip.modules import module_dict, module_params_dict

# from .lcat import *
module_dict_gnn = {}
module_dict_gnn['gcn'] = GCN
module_dict_gnn['gat'] = GAT
module_dict_gnn['gin'] = GIN
module_dict_gnn['gine'] = GINE
module_dict_gnn['pna'] = PNA
module_dict_gnn["diffpool"] = DiffPool
# module_dict_gnn['lcatgcn'] = LCATGCN

module_dict.update(module_dict_gnn)
for module_name, module_class in module_dict_gnn.items():
    module_params_dict[module_name] = module_class.kwargs

pooling_dict = {}
