[Pytorch Geometric] Convert dense adjacency matrix to sparse matrix

Hi, i want to convert a batched dense edge adjacency matrix of size (B,N,N) to a batched sparse edge adjacency matrix of size (2, M), in which B denotes the batch size, N denotes the maximum number of nodes each graph and M denotes the number of edges in one batch.

I could only find one function for this purpose in the package torch_geometric.utils named dense_to_sparse. However, the source code shows that this function does not work for batched data. My solution is to iterate over batch, apply this function on each graph, and aggregate the results to get the batched result. I wonder if this iteration would impact the calculation performance and is there any better way to achieve the same purpose.

The other question is that when i convert a batched sparse edge adjacency matrix of size (2,M) with edge attributes of size (M, F), in which F denotes the dimension of features for each edge, to a batched dense edge adjacency matrix using the function to_dense_adj provided by package torch_geometric.utils, It results a tensor of size (B, N, N, F), but I couldn’t find the function for converting such a tensor back. Is there any convenient way to convert such a tensor back to sparse edge adjacency matrix of size (2,M) and edge attributes of size (M, F)?

For anyone who have the same problem, here is the solution.
torch_geometric discussion on github