fix #3786

pyg-team · Jan 5, 2022 · 690a01d · 690a01d
1 parent 0ce9bf6
commit 690a01d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_geometric/nn/conv/appnp.py b/torch_geometric/nn/conv/appnp.py
@@ -121,8 +121,8 @@ def forward(self, x: Tensor, edge_index: Adj,
 
         return x
 
-    def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
-        return edge_weight.view(-1, 1) * x_j
+    def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
+        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j
 
     def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
         return matmul(adj_t, x, reduce=self.aggr)