TS-EGNN doesn't work smoothly with PyG 2.5

[xiaoruiDong]

Haven't figured out the reason yet. Potentially related to this issue.

TS-EGNN includes a modified version of EGNN implementation from https://github.com/lucidrains/egnn-pytorch. Even after incorporating the new changes of PyG (changing attribute names and so on) to the TS-EGNN code, TS-EGNN is still not working.

The error is raised at

hidden_out, coors_out = self.propagate(edge_index, x=feats, edge_attr=edge_attr_feats,
                                                           coors=coors, rel_coors=rel_coors, 
                                                           batch=batch)

# where the last step of propagate is 
return self.update((hidden_out, coors_out), **update_kwargs)

and the error implies the returned item is not a tuple.

But very weirdly, when I inserted a pdb.set_trace() inside the propagate and simply "continue"d it after stopping, everything worked. The same script also works on PyG 2.4.0 without issues. Inspired by the issue above, it makes me wonder if this relates to dynamic compilation or something similar, especially given propagate is re-implemented in the egnn code. However, no further action is planned at the moment

Current solution: use PyTorch 2.1.2 + PyG 2.4.0 instead