(self, g: Graph)
| 8 | self.pn = pn |
| 9 | |
| 10 | def augment(self, g: Graph) -> Graph: |
| 11 | x, edge_index, edge_weights = g.unfold() |
| 12 | |
| 13 | edge_index, edge_weights = drop_node(edge_index, edge_weights, keep_prob=1. - self.pn) |
| 14 | |
| 15 | return Graph(x=x, edge_index=edge_index, edge_weights=edge_weights) |