The code on torch-geometric 1.5.0 fails in 2.1.0

[WeixuanXiong]

The false is as followed:
Traceback (most recent call last):
File "model/graph_learning.py", line 57, in
out, adj = model(x)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "model/graph_learning.py", line 49, in forward
gcn_out = gnn_layers(x, batch_all_edge_index, node_num=node_num * batch_num, embedding=all_embeddings)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/dfs/data/Anomaly-Transformer/model/graph_layer.py", line 146, in forward
out, (new_edge_index, att_weight) = self.gnn(x, edge_index, embedding, return_attention_weights=True)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/dfs/data/Anomaly-Transformer/model/graph_layer.py", line 66, in forward
return_attention_weights=return_attention_weights)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_geometric/nn/conv/message_passing.py", line 391, in propagate
out = self.aggregate(out, **aggr_kwargs)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_geometric/nn/conv/message_passing.py", line 515, in aggregate
dim=self.node_dim)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_geometric/nn/aggr/base.py", line 114, in call
return super().call(x, index, ptr, dim_size, dim, **kwargs)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_geometric/nn/aggr/basic.py", line 21, in forward
return self.reduce(x, index, ptr, dim_size, dim, reduce='sum')
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_geometric/nn/aggr/base.py", line 153, in reduce
return scatter(x, index, dim=dim, dim_size=dim_size, reduce=reduce)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_geometric/utils/scatter.py", line 64, in scatter
return torch_scatter.scatter(src, index, dim, out, dim_size, reduce)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_scatter/scatter.py", line 152, in scatter
return scatter_sum(src, index, dim, out, dim_size)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_scatter/scatter.py", line 11, in scatter_sum
index = broadcast(index, src, dim)
File "/dfs/data/miniconda/envs/anod/lib/python3.7/site-packages/torch_scatter/utils.py", line 12, in broadcast
src = src.expand(other.size())
RuntimeError: The expanded size of the tensor (1) must match the existing size (160000) at non-singleton dimension 1. Target sizes: [160000, 1, 64]. Tensor sizes: [1, 160000, 1]

The MessagePassing class is:

class GraphLayer(MessagePassing):
def init(self, in_channels, out_channels, heads=1, concat=True,
negative_slope=0.2, dropout=0, bias=True, inter_dim=-1,**kwargs):
super(GraphLayer, self).init(aggr='add', **kwargs)

    self.in_channels = in_channels
    self.out_channels = out_channels
    self.heads = heads
    self.concat = concat
    self.negative_slope = negative_slope
    self.dropout = dropout

    self.__alpha__ = None

    self.lin = Linear(in_channels, heads * out_channels, bias=False)

    self.att_i = Parameter(torch.Tensor(1, heads, out_channels))
    self.att_j = Parameter(torch.Tensor(1, heads, out_channels))
    self.att_em_i = Parameter(torch.Tensor(1, heads, out_channels))
    self.att_em_j = Parameter(torch.Tensor(1, heads, out_channels))

    if bias and concat:
        self.bias = Parameter(torch.Tensor(heads * out_channels))
    elif bias and not concat:
        self.bias = Parameter(torch.Tensor(out_channels))
    else:
        self.register_parameter('bias', None)

    self.reset_parameters()

def reset_parameters(self):
    glorot(self.lin.weight)
    glorot(self.att_i)
    glorot(self.att_j)
    
    zeros(self.att_em_i)
    zeros(self.att_em_j)

    zeros(self.bias)



def forward(self, x, edge_index, embedding, return_attention_weights=False):
    """"""
    if torch.is_tensor(x):
        x = self.lin(x)
        x = (x, x)
    else:
        x = (self.lin(x[0]), self.lin(x[1]))

    edge_index, _ = remove_self_loops(edge_index)
    edge_index, _ = add_self_loops(edge_index,
                                   num_nodes=x[1].size(self.node_dim))

    out = self.propagate(edge_index, x=x, embedding=embedding, edges=edge_index,
                         return_attention_weights=return_attention_weights)

    if self.concat:
        out = out.view(-1, self.heads * self.out_channels)
    else:
        out = out.mean(dim=1)

    if self.bias is not None:
        out = out + self.bias

    if return_attention_weights:
        alpha, self.__alpha__ = self.__alpha__, None
        return out, (edge_index, alpha)
    else:
        return out

def message(self, x_i, x_j, edge_index_i, size_i,
            embedding,
            edges,
            return_attention_weights):
    x_i = x_i.view(-1, self.heads, self.out_channels)
    x_j = x_j.view(-1, self.heads, self.out_channels)
    if embedding is not None:
        embedding_i, embedding_j = embedding[edge_index_i], embedding[edges[0]]
        embedding_i = embedding_i.unsqueeze(1).repeat(1,self.heads,1)
        embedding_j = embedding_j.unsqueeze(1).repeat(1,self.heads,1)

        key_i = torch.cat((x_i, embedding_i), dim=-1)
        key_j = torch.cat((x_j, embedding_j), dim=-1)

    cat_att_i = torch.cat((self.att_i, self.att_em_i), dim=-1)
    cat_att_j = torch.cat((self.att_j, self.att_em_j), dim=-1)

    alpha = (key_i * cat_att_i).sum(-1) + (key_j * cat_att_j).sum(-1)

    alpha = alpha.view(-1, self.heads, 1)


    alpha = F.leaky_relu(alpha, self.negative_slope)
    alpha = softmax(alpha, edge_index_i, num_nodes=size_i)

    if return_attention_weights:
        self.__alpha__ = alpha

    alpha = F.dropout(alpha, p=self.dropout, training=self.training)
    print(x_j.shape)
    print(alpha.shape)
    return x_j * alpha.view(-1, self.heads, 1)

The result of last two prints are: (16000, 1, 64) and (16000, 1, 1).
This class works well on torch 1.5.0 with torch-geometric 1.5.0 but fails on torch 1.12.0 with torch-geometric 2.1.0. I'm wondering how can i get rid of this.

[EdisonLeeeee]

You can simply try

        return (x_j * alpha.view(-1, self.heads, 1)).squeeze()

[Beck-Pro]

YOU SAVED MY LIFE THANK YOU Edison!!