Compatibility issues with torch_scatter

[aasomol]

Hello,
I need to reproduce a code which is a bit old and uses the following versions (the arrow points to what I have installed right now) because of some compatibility issues with my GPU:

torch-cluster==1.4.5 ---> 1.6.0
torch-geometric==1.3.2 ---> 2.0.4
torch-scatter==1.4.0 ---> 2.0.9
torch-sparse==0.4.3 ---> 0.6.13
torch-spline-conv==1.1.0 ---> 1.2.1

My question is whether you know if the function torch_geometric.utils.softmax has experienced a major change in code, because it works with the old version but now with the new one. It seems the problem is in torch-scatter:

pi = U.softmax(pi.squeeze(),state_batch.glob_action_batch)

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
File "/.../python3.9/site-packages/torch_geometric/utils/softmax.py", line 41, in softmax
elif index is not None:
N = maybe_num_nodes(index, num_nodes)
src_max = scatter(src, index, dim, dim_size=N, reduce='max')
~~~~~~~ <--- HERE
src_max = src_max.index_select(dim, index)
out = (src - src_max).exp()
File "/.../python3.9/site-packages/torch_scatter/scatter.py", line 160, in scatter
return scatter_min(src, index, dim, out, dim_size)[0]
elif reduce == 'max':
return scatter_max(src, index, dim, out, dim_size)[0]
~~~~~~~~~~~ <--- HERE
else:
raise ValueError
File "/.../python3.9/site-packages/torch_scatter/scatter.py", line 72, in scatter_max
out: Optional[torch.Tensor] = None,
dim_size: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
return torch.ops.torch_scatter.scatter_max(src, index, dim, out, dim_size)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
RuntimeError: The expanded size of the tensor (126241) must match the existing size (851) at non-singleton dimension 0. Target sizes: [126241]. Tensor sizes: [851]

Thank you so much in advance.

[rusty1s]

The softmax code looks nearly identical to me in 1.3.2 and 2.0.4, so I assume the real "bug" happens somewhere else. What are the shapes of pi.squeeze() and state_batch.glob_action_batch? The sizes of these tensors have to match in the first dimension.