APPNP applies edge dropout cumulatively across propagation steps

Describe the bug

I believe there is a bug in the forward pass of APPNP (torch_geometric/nn/conv/appnp.py).

In the current behaviour: https://github.com/pyg-team/pytorch_geometric/blob/c7ccd0751564153dcffcbff03aea2829fa8db57e/torch_geometric/nn/conv/appnp.py#L121

we have edge_weight = F.dropout(edge_weight, p=self.dropout). This means that edge dropout is applied cumulatively across the K propagation steps, rather than independently at each step. This causes edge weights to vanish much faster than intended, effectively removing most edges after only a few iterations. This also means that weights of surviving edges increase due to the scaling factor (1/(1-p)) applied during dropout, which can blow up and cause NaNs.

Expected behaviour for the block where isinstance(edge_index, Tensor) and not is_torch_sparse_tensor(edge_index)

        original_edge_weight = edge_weight
        for _ in range(self.K):
            if self.dropout > 0 and self.training:
                if isinstance(edge_index, Tensor):
                    if is_torch_sparse_tensor(edge_index):
                        _, edge_weight = to_edge_index(edge_index)
                        edge_weight = F.dropout(edge_weight, p=self.dropout)
                        edge_index = set_sparse_value(edge_index, edge_weight)
                    else:
                        assert edge_weight is not None
                        edge_weight = F.dropout(original_edge_weight, p=self.dropout)
                else:
                    value = edge_index.storage.value()
                    assert value is not None
                    value = F.dropout(value, p=self.dropout)
                    edge_index = edge_index.set_value(value, layout='coo')
            x = self.propagate(edge_index, x=x, edge_weight=edge_weight)
            x = x * (1 - self.alpha)
            x = x + self.alpha * h

In the parts where edge_index is used to generate the weights I'm not sure on the fix.

Versions

Versions of relevant libraries: [pip3] numpy==2.3.3 [pip3] nvidia-cublas-cu11==11.11.3.6 [pip3] nvidia-cuda-cupti-cu11==11.8.87 [pip3] nvidia-cuda-nvrtc-cu11==11.8.89 [pip3] nvidia-cuda-runtime-cu11==11.8.89 [pip3] nvidia-cudnn-cu11==9.1.0.70 [pip3] nvidia-cufft-cu11==10.9.0.58 [pip3] nvidia-curand-cu11==10.3.0.86 [pip3] nvidia-cusolver-cu11==11.4.1.48 [pip3] nvidia-cusparse-cu11==11.7.5.86 [pip3] nvidia-nccl-cu11==2.21.5 [pip3] nvidia-nccl-cu13==2.31.2 [pip3] nvidia-nvtx-cu11==11.8.86 [pip3] pytorch-lightning==2.6.0 [pip3] torch==2.6.0+cu118 [pip3] torch-geometric==2.7.0 [pip3] torchaudio==2.6.0+cu118 [pip3] torchmetrics==1.8.2 [pip3] torchvision==0.21.0+cu118 [pip3] triton==3.2.0 [conda] Could not collect