Memory leak in GraphUNet due to function to_torch_csr_tensor

[acotino-ignitioncomputing]

🐛 Describe the bug

Consecutive application of the forward-pass of GraphUNet causes a steady increase of memory, which eventually leads to a Memory Overflow. Changing the function to_torch_csr_tensor inside the method GraphUNet.augment_adj to the function to_torch_coo_tensor seems to solve this issue.
Note: I was not able to reproduce this steady increase of memory by application of to_torch_csr_tensor alone.

import torch
from torch_geometric.nn import GraphUNet

if __name__ == "__main__":
    # Initialize GraphUNet-model
    model = GraphUNet(3, 4, 1, depth=4)

    # Define sufficient large input tensors to make impact of memory leak visible
    N = 9000
    feature_matrix = torch.rand([N, 3])
    indices = torch.range(0, N - 1, dtype=int)
    edge_index = torch.stack((indices, indices))

    cycles = 500
    prompt = ""
    while not prompt:
        for _ in range(cycles):
            model(feature_matrix, edge_index)
        print(f"Performed {cycles} forward-calls")
        prompt = input("Provide empty input to continue: ")

Versions

PyG version: 2.7.0
PyTorch version: 2.4.1+cpu
Python version: 3.10.12