update

test/nn/models/test_basic_gnn.py

@@ -161,6 +161,7 @@ def test_packaging():
     path = osp.join(torch.hub._get_torch_home(), 'pyg_test_package.pt')
     with torch.package.PackageExporter(path) as pe:
         pe.extern('torch_geometric.nn.**')
+        pe.extern('torch_geometric.utils.trim_to_layer')
         pe.extern('_operator')
         pe.save_pickle('models', 'model.pkl', model)
 

torch_geometric/nn/models/basic_gnn.py

@@ -141,7 +141,9 @@ def __init__(
                 in_channels = hidden_channels
             self.lin = Linear(in_channels, self.out_channels)
 
-        self.trim = TrimToLayer()
+        # We define `trim_to_layer` functionality as a module such that we can
+        # still use `to_hetero` on-top.
+        self._trim = TrimToLayer()
 
     def init_conv(self, in_channels: Union[int, Tuple[int, int]],
                   out_channels: int, **kwargs) -> MessagePassing:
@@ -187,12 +189,6 @@ def forward(
                 scenarios to only operate on minimal-sized representations.
                 (default: :obj:`None`)
         """
-        if (num_sampled_edges_per_hop is not None
-                and num_sampled_nodes_per_hop is None):
-            raise ValueError("'num_sampled_nodes_per_hop' needs to be given")
-        if (num_sampled_nodes_per_hop is not None
-                and num_sampled_edges_per_hop is None):
-            raise ValueError("'num_sampled_edges_per_hop' needs to be given")
         if (num_sampled_nodes_per_hop is not None
                 and isinstance(edge_weight, Tensor)
                 and isinstance(edge_attr, Tensor)):
@@ -202,8 +198,8 @@ def forward(
 
         xs: List[Tensor] = []
         for i in range(self.num_layers):
-            if num_sampled_nodes_per_hop is not None:
-                x, edge_index, value = self.trim(
+            if isinstance(num_sampled_nodes_per_hop, (list, tuple)):
+                x, edge_index, value = self._trim(
                     i,
                     num_sampled_nodes_per_hop,
                     num_sampled_edges_per_hop,

torch_geometric/nn/summary.py

@@ -77,6 +77,9 @@ def hook(module, inputs, output):
         name, module, depth = stack.pop()
         module_id = id(module)
 
+        if name.startswith('(_'):  # Do not summarize private modules.
+            continue
+
         if module_id in hooks:  # Avoid duplicated hooks.
             hooks[module_id].remove()
 

torch_geometric/utils/trim_to_layer.py

@@ -103,6 +103,13 @@ def forward(
         edge_attr: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
 
+        if (num_sampled_nodes_per_hop is None
+                and num_sampled_edges_per_hop is not None):
+            raise ValueError("'num_sampled_nodes_per_hop' needs to be given")
+        if (num_sampled_edges_per_hop is None
+                and num_sampled_nodes_per_hop is not None):
+            raise ValueError("'num_sampled_edges_per_hop' needs to be given")
+
         if num_sampled_nodes_per_hop is None:
             return x, edge_index, edge_attr
         if num_sampled_edges_per_hop is None: