1f0750670cf8ea24ad264debf9603002ab0fb565,torch_geometric/nn/pool/mem_pool.py,MemPool,forward,#MemPool#Any#Any#,77

Before Change

        B = x.shape[0]
        H = self.heads
        K = self.num_keys
        N = x.shape[1]

        dist = torch.cdist(self.k.view(H * K, -1), x.view(B * N, -1), 1)
        dist = (1 + (dist**2 / self.tau)).pow(-(self.tau + 1) / 2.0)
        dist = dist.view(H, K, B, N).permute(2, 0, 3, 1)
        S = dist / dist.sum(3).unsqueeze(3)  // B*H*N*K

        S = self.conv(S).squeeze(1)
        S = torch.softmax(S, 2)  // B*N*K
        S = self.apply_mask(S, mask)
        x = self.lin(matmul(S.transpose(1, 2), x))  // B*K*F"
        x = self.act(x)

After Change

        
        x, mask = to_dense_batch(x, batch)

        (B, N, _), H, K = x.size(), self.heads, self.num_clusters

        dist = torch.cdist(self.k.view(H * K, -1), x.view(B * N, -1), p=2)**2
        dist = (1. + dist / self.tau).pow(-(self.tau + 1.0) / 2.0)

Instances