TypeError when fitting GADNR model: MessagePassing.__init__() got an unexpected keyword argument 'tot_nodes'

[alenkaja00]

Describe the bug
When running model.fit with a GADNR model, GADNRBase is called with a parameter tot_nodes which MessagePassing layer does not accept.

To Reproduce
Steps to reproduce the behavior:

1. Instantiate model = GADNR(gpu=0, hid_dim=16, num_layers=2, batch_size=4096, num_neigh=[5,5], epoch=50)
2. Error will appar TypeError: MessagePassing.__init__() got an unexpected keyword argument 'tot_nodes'

Expected behavior
Model fits without errors.

Screenshots
Complete error below:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File <timed eval>:1

File ~/.local/lib/python3.11/site-packages/pygod/detector/gadnr.py:311, in GADNR.fit(self, data, label, h_loss_weight, degree_loss_weight, feature_loss_weight, loss_step)
    307     loader = NeighborLoader(data,
    308                             self.num_neigh,
    309                             batch_size=self.batch_size)
    310     self.full_batch = False
--> 311 self.model = self.init_model(**self.kwargs)
    312 if self.compile_model:
    313     self.model = compile(self.model)

File ~/.local/lib/python3.11/site-packages/pygod/detector/gadnr.py:204, in GADNR.init_model(self, **kwargs)
    201 if self.save_emb:
    202     self.emb = torch.zeros(self.num_nodes, self.hid_dim)
--> 204 return GADNRBase(in_dim=self.in_dim, hid_dim=self.hid_dim,
    205                  encoder_layers=self.encoder_layers,
    206                  deg_dec_layers=self.deg_dec_layers,
    207                  fea_dec_layers=self.fea_dec_layers,
    208                  sample_size=self.sample_size,
    209                  sample_time=self.sample_time, 
    210                  neighbor_num_list=self.neighbor_num_list,
    211                  tot_nodes=self.tot_nodes,
    212                  neigh_loss=self.neigh_loss,
    213                  lambda_loss1=self.lambda_loss1,
    214                  lambda_loss2=self.lambda_loss2,
    215                  lambda_loss3=self.lambda_loss3,
    216                  full_batch=self.full_batch,
    217                  backbone=self.backbone,
    218                  device=self.device).to(self.device)

File ~/.local/lib/python3.11/site-packages/pygod/nn/gadnr.py:136, in GADNRBase.__init__(self, in_dim, hid_dim, encoder_layers, deg_dec_layers, fea_dec_layers, sample_size, sample_time, neighbor_num_list, neigh_loss, lambda_loss1, lambda_loss2, lambda_loss3, full_batch, dropout, act, backbone, device, **kwargs)
    133 self.mlp_gen = MLP_generator(hid_dim, hid_dim)
    135 # Encoder
--> 136 self.shared_encoder = backbone(in_channels=hid_dim,
    137                                hidden_channels=hid_dim,
    138                                num_layers=encoder_layers,
    139                                out_channels=hid_dim,
    140                                dropout=dropout,
    141                                act=act,
    142                                **kwargs)
    144 # Decoder
    145 self.degree_decoder = FNN_GAD_NR(hid_dim, hid_dim, 1, deg_dec_layers)

File ~/.local/lib/python3.11/site-packages/torch_geometric/nn/models/basic_gnn.py:106, in BasicGNN.__init__(self, in_channels, hidden_channels, num_layers, out_channels, dropout, act, act_first, act_kwargs, norm, norm_kwargs, jk, **kwargs)
    103 self.convs = ModuleList()
    104 if num_layers > 1:
    105     self.convs.append(
--> 106         self.init_conv(in_channels, hidden_channels, **kwargs))
    107     if isinstance(in_channels, (tuple, list)):
    108         in_channels = (hidden_channels, hidden_channels)

File ~/.local/lib/python3.11/site-packages/torch_geometric/nn/models/basic_gnn.py:430, in GCN.init_conv(self, in_channels, out_channels, **kwargs)
    428 def init_conv(self, in_channels: int, out_channels: int,
    429               **kwargs) -> MessagePassing:
--> 430     return GCNConv(in_channels, out_channels, **kwargs)

File ~/.local/lib/python3.11/site-packages/torch_geometric/nn/conv/gcn_conv.py:190, in GCNConv.__init__(self, in_channels, out_channels, improved, cached, add_self_loops, normalize, bias, **kwargs)
    178 def __init__(
    179     self,
    180     in_channels: int,
   (...)
    187     **kwargs,
    188 ):
    189     kwargs.setdefault('aggr', 'add')
--> 190     super().__init__(**kwargs)
    192     if add_self_loops is None:
    193         add_self_loops = normalize

TypeError: MessagePassing.__init__() got an unexpected keyword argument 'tot_nodes'

Desktop (please complete the following information):

• OS: Linux
• Environment JupyterHub

[kayzliu]

@YingtongDou Could you please help with this issue?