fixing tests

pyg-team · Jun 3, 2023 · 2b576bb · 2b576bb
1 parent 352df4b
commit 2b576bb
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/test/nn/test_encoding.py b/test/nn/test_encoding.py
@@ -1,12 +1,11 @@
 import torch
 
-from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import (
     NodeEncoding,
     PositionalEncoding,
     TemporalEncoding,
 )
-from torch_geometric.testing import withCUDA
+from torch_geometric.testing import onlyNeighborSampler, withCUDA
 
 
 @withCUDA
@@ -28,7 +27,9 @@ def test_temporal_encoding(device):
 
 
 @withCUDA
+@onlyNeighborSampler
 def test_node_encoding(get_dataset, device):
+    from torch_geometric.loader import NeighborLoader
     dataset = get_dataset(name='Cora')
     data = dataset[0]
     loader = NeighborLoader(

diff --git a/torch_geometric/nn/encoding.py b/torch_geometric/nn/encoding.py
@@ -103,17 +103,17 @@ class NodeEncoding(torch.nn.Module):
     Model Architectures for Temporal Networks?"
     <https://openreview.net/forum?id=ayPPc0SyLv1>`_ paper.
     :class:`NodeEncoding` captures the node identity and node feature
-        information via neighbor mean-pooling
+    information via neighbor mean-pooling
 
     """
-    def forward(self, data):
+    def forward(self, data) -> Tensor:
         """
-
         Args:
             data: batch of nodes from NeighborLoader
 
         Returns:
-
+            torch.Tensor of the root nodes updated with mean
+                pooling of neighbor features
         """
         x, edge_index, batch_size = data.x, data.edge_index, data.batch_size
         root_nodes = x[:, batch_size]