fix: RuntimeError in GraphIsomorphismNetwork #92 (#93)

Author: kanojikajino

torchdrug/models/chebnet.py

@@ -35,7 +35,7 @@ def __init__(self, input_dim, hidden_dims, edge_input_dim=None, k=1, short_cut=F
         if not isinstance(hidden_dims, Sequence):
             hidden_dims = [hidden_dims]
         self.input_dim = input_dim
-        self.output_dim = hidden_dims[-1] * (len(hidden_dims) if concat_hidden else 1)
+        self.output_dim = sum(hidden_dims) if concat_hidden else hidden_dims[-1]
         self.dims = [input_dim] + list(hidden_dims)
         self.short_cut = short_cut
         self.concat_hidden = concat_hidden

torchdrug/models/gat.py

@@ -35,7 +35,7 @@ def __init__(self, input_dim, hidden_dims, edge_input_dim=None, num_head=1, nega
         if not isinstance(hidden_dims, Sequence):
             hidden_dims = [hidden_dims]
         self.input_dim = input_dim
-        self.output_dim = hidden_dims[-1] * (len(hidden_dims) if concat_hidden else 1)
+        self.output_dim = sum(hidden_dims) if concat_hidden else hidden_dims[-1]
         self.dims = [input_dim] + list(hidden_dims)
         self.short_cut = short_cut
         self.concat_hidden = concat_hidden

torchdrug/models/gcn.py

@@ -33,7 +33,7 @@ def __init__(self, input_dim, hidden_dims, edge_input_dim=None, short_cut=False,
         if not isinstance(hidden_dims, Sequence):
             hidden_dims = [hidden_dims]
         self.input_dim = input_dim
-        self.output_dim = hidden_dims[-1] * (len(hidden_dims) if concat_hidden else 1)
+        self.output_dim = sum(hidden_dims) if concat_hidden else hidden_dims[-1]
         self.dims = [input_dim] + list(hidden_dims)
         self.short_cut = short_cut
         self.concat_hidden = concat_hidden

torchdrug/models/gin.py

@@ -37,7 +37,7 @@ def __init__(self, input_dim=None, hidden_dims=None, edge_input_dim=None, num_ml
         if not isinstance(hidden_dims, Sequence):
             hidden_dims = [hidden_dims]
         self.input_dim = input_dim
-        self.output_dim = hidden_dims[-1] * (len(hidden_dims) if concat_hidden else 1)
+        self.output_dim = sum(hidden_dims) if concat_hidden else hidden_dims[-1]
         self.dims = [input_dim] + list(hidden_dims)
         self.short_cut = short_cut
         self.concat_hidden = concat_hidden