Masking implementation confusion

[yuewan2]

In the paper, it mentions that you apply random masking on atom and bonds. It also says that "GROVER randomly masks a local subgraph". However, in the GroverCollator class, the 'atom_random_mask' and 'bond_random_mask' are only designed for selecting atom/bond context without applying any masking on the input molecule representation. I am wondering if there is anything I missed from the code since I do not see any molecule masking implementation elsewhere.