RuntimeError in GraphIsomorphismNetwork

[kanojikajino]

Hi,

I was playing with the property prediction tutorial (https://torchdrug.ai/docs/tutorials/property_prediction.html), and encountered the following RuntimeError, if I tweak hidden_dims as,

model = models.GIN(input_dim=dataset.node_feature_dim,
                   hidden_dims=[128, 192],
                   short_cut=True, batch_norm=True, concat_hidden=True)

and the error message says:

Traceback (most recent call last):
  File "/home/docker/tasks/mount/torchdrug_playground/test.py", line 20, in <module>
    solver.train(num_epoch=200)
  File "/home/docker/torchdrug/torchdrug/core/engine.py", line 155, in train
    loss, metric = model(batch)
  File "/home/docker/.pyenv/versions/3.9.0/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/docker/torchdrug/torchdrug/tasks/property_prediction.py", line 74, in forward
    pred = self.predict(batch, all_loss, metric)
  File "/home/docker/torchdrug/torchdrug/tasks/property_prediction.py", line 106, in predict
    pred = self.linear(output["graph_feature"])
  File "/home/docker/.pyenv/versions/3.9.0/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/docker/.pyenv/versions/3.9.0/lib/python3.9/site-packages/torch/nn/modules/linear.py", line 103, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1024x320 and 384x2)

It seems that the output dimension of the model should be

self.output_dim = sum(hidden_dims) if concat_hidden else hidden_dims[-1]

instead of the current implementation:

torchdrug/torchdrug/models/gin.py

Line 40 in 26f15f6

self.output_dim = hidden_dims[-1] * (len(hidden_dims) if concat_hidden else 1)

If the correction above is ok, I will create a pull request to fix it. Or, if there is any misunderstanding, I would appreciate if you could point it out.

[Oxer11]

Hi! Thanks for your issue.

Yes, you're right. This is actually a bug. The output dim should be the sum of all hidden dims if concat_hidden is True. Thanks for point out. Could you create a pull request to fix it?

[kanojikajino]

Hi, I created a pull request to fix output_dim issues in several modules (#93). I would appreciate if you could review it.