I am talking about this official implementation: https://github.com/vgsatorras/egnn/blob/main/models/egnn_clean/egnn_clean.py