How to set DropEdge for SparseTensor?

[realCrush]

Hi, I want to use SparseTensor type adjacent matrix for reproducibility in message passing, however when I want to set DropEdge for it, it seems wired, my code:

        row, col, edge_attr_ = adj_t.coo() # adj_t: SparseTensor
        edge_index_ = torch.stack([row, col], dim=0)
        edge_index_, edge_attr_ = dropout_adj(training=self.training)
        adj_t = SparseTensor.from_edge_index(edge_index_, edge_attr_).t()

above code will trigger error like:

  File "message_passing.py", line 202, in __collect__
    self.__set_size__(size, dim, data)
  File "message_passing.py.py", line 165, in __set_size__
    raise ValueError(
ValueError: Encountered tensor with size 5201 in dimension -2, but expected size 5200.

What should I do to do DropEdge on SparseTensor? I would be super appreciate if anyone can help me 🙏🏻

[rusty1s]

You should also pass in the sparse_size of the adjacency matrix in the SparseTensor.from_edge_index call to ensure that the size is preserved. Also happy to take a PR that directly integrates this functionality into dropout_adj.

[realCrush]

Thank you rusty1s, you are my superman 🦸🏻‍♂️