egnn Just wanted to make clear that this implementation is not complete. I started working on it, but ran into some issues with PyTorch's sparse tensor support, I don't remember exactly what.