pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/distributed/local_feature_store.py

@@ -15,6 +15,7 @@ from torch_geometric.distributed.rpc import (
     rpc_async,
     rpc_register,
 )
+from torch_geometric.io import fs
 from torch_geometric.typing import EdgeType, NodeOrEdgeType, NodeType
 
 
@@ -415,11 +416,11 @@ class LocalFeatureStore(FeatureStore):
 
         node_feats: Optional[Dict[str, Any]] = None
         if osp.exists(osp.join(part_dir, 'node_feats.pt')):
-            node_feats = torch.load(osp.join(part_dir, 'node_feats.pt'))
+            node_feats = fs.torch_load(osp.join(part_dir, 'node_feats.pt'))
 
         edge_feats: Optional[Dict[str, Any]] = None
         if osp.exists(osp.join(part_dir, 'edge_feats.pt')):
-            edge_feats = torch.load(osp.join(part_dir, 'edge_feats.pt'))
+            edge_feats = fs.torch_load(osp.join(part_dir, 'edge_feats.pt'))
 
         if not meta['is_hetero'] and node_feats is not None:
             feat_store.put_global_id(node_feats['global_id'], group_name=None)

torch_geometric/distributed/local_graph_store.py

@@ -6,6 +6,7 @@ from torch import Tensor
 
 from torch_geometric.data import EdgeAttr, GraphStore
 from torch_geometric.distributed.partition import load_partition_info
+from torch_geometric.io import fs
 from torch_geometric.typing import EdgeTensorType, EdgeType, NodeType
 from torch_geometric.utils import sort_edge_index
 
@@ -185,7 +186,7 @@ class LocalGraphStore(GraphStore):
         graph_store.edge_pb = edge_pb
         graph_store.meta = meta
 
-        graph_data = torch.load(osp.join(part_dir, 'graph.pt'))
+        graph_data = fs.torch_load(osp.join(part_dir, 'graph.pt'))
         graph_store.is_sorted = meta['is_sorted']
 
         if not meta['is_hetero']:

torch_geometric/distributed/partition.py

@@ -3,15 +3,16 @@ import logging
 import os
 import os.path as osp
 from collections import defaultdict
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 
 import torch_geometric.distributed as pyg_dist
 from torch_geometric.data import Data, HeteroData
+from torch_geometric.io import fs
 from torch_geometric.loader.cluster import ClusterData
 from torch_geometric.sampler.utils import sort_csc
-from torch_geometric.typing import Dict, EdgeType, EdgeTypeStr, NodeType, Tuple
+from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType
 
 
 class Partitioner:
@@ -303,7 +304,7 @@ class Partitioner:
                 elif self.is_node_level_time:
                     node_time = data.time
 
-                # Sort by column to avoid keeping track of permuations in
+                # Sort by column to avoid keeping track of permutations in
                 # `NeighborSampler` when converting to CSC format:
                 global_row, global_col, perm = sort_csc(
                     global_row, global_col, node_time, edge_time)
@@ -360,7 +361,7 @@ class Partitioner:
             'edge_types': self.edge_types,
             'node_offset': list(node_offset.values()) if node_offset else None,
             'is_hetero': self.is_hetero,
-            'is_sorted': True,  # Based on colum/destination.
+            'is_sorted': True,  # Based on column/destination.
         }
         with open(osp.join(self.root, 'META.json'), 'w') as f:
             json.dump(meta, f)
@@ -380,21 +381,21 @@ def load_partition_info(
     assert osp.exists(partition_dir)
 
     if meta['is_hetero'] is False:
-        node_pb = torch.load(osp.join(root_dir, 'node_map.pt'))
-        edge_pb = torch.load(osp.join(root_dir, 'edge_map.pt'))
+        node_pb = fs.torch_load(osp.join(root_dir, 'node_map.pt'))
+        edge_pb = fs.torch_load(osp.join(root_dir, 'edge_map.pt'))
 
         return (meta, num_partitions, partition_idx, node_pb, edge_pb)
     else:
         node_pb_dict = {}
         node_pb_dir = osp.join(root_dir, 'node_map')
         for ntype in meta['node_types']:
-            node_pb_dict[ntype] = torch.load(
+            node_pb_dict[ntype] = fs.torch_load(
                 osp.join(node_pb_dir, f'{pyg_dist.utils.as_str(ntype)}.pt'))
 
         edge_pb_dict = {}
         edge_pb_dir = osp.join(root_dir, 'edge_map')
         for etype in meta['edge_types']:
-            edge_pb_dict[tuple(etype)] = torch.load(
+            edge_pb_dict[tuple(etype)] = fs.torch_load(
                 osp.join(edge_pb_dir, f'{pyg_dist.utils.as_str(etype)}.pt'))
 
         return (meta, num_partitions, partition_idx, node_pb_dict,

torch_geometric/distributed/rpc.py

@@ -92,7 +92,7 @@ def shutdown_rpc(id: str = None, graceful: bool = True,
 class RPCRouter:
     r"""A router to get the worker based on the partition ID."""
     def __init__(self, partition_to_workers: List[List[str]]):
-        for pid, rpc_worker_list in enumerate(partition_to_workers):
+        for rpc_worker_list in partition_to_workers:
             if len(rpc_worker_list) == 0:
                 raise ValueError('No RPC worker is in worker list')
         self.partition_to_workers = partition_to_workers
@@ -120,7 +120,7 @@ def rpc_partition_to_workers(
     partition_to_workers = [[] for _ in range(num_partitions)]
     gathered_results = global_all_gather(
         (ctx.role, num_partitions, current_partition_idx))
-    for worker_name, (role, nparts, idx) in gathered_results.items():
+    for worker_name, (_, _, idx) in gathered_results.items():
         partition_to_workers[idx].append(worker_name)
     return partition_to_workers
 
@@ -144,7 +144,7 @@ _rpc_call_pool: Dict[int, RPCCallBase] = {}
 @rpc_require_initialized
 def rpc_register(call: RPCCallBase) -> int:
     r"""Registers a call for RPC requests."""
-    global _rpc_call_id, _rpc_call_pool
+    global _rpc_call_id
 
     with _rpc_call_lock:
         call_id = _rpc_call_id

torch_geometric/edge_index.py

@@ -17,6 +17,7 @@ from typing import (
     overload,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor
@@ -173,7 +174,7 @@ class EdgeIndex(Tensor):
     :meth:`EdgeIndex.fill_cache_`, and are maintained and adjusted over its
     lifespan (*e.g.*, when calling :meth:`EdgeIndex.flip`).
 
-    This representation ensures for optimal computation in GNN message passing
+    This representation ensures optimal computation in GNN message passing
     schemes, while preserving the ease-of-use of regular COO-based :pyg:`PyG`
     workflows.
 
@@ -183,7 +184,7 @@ class EdgeIndex(Tensor):
 
         edge_index = EdgeIndex(
             [[0, 1, 1, 2],
-             [1, 0, 2, 1]]
+             [1, 0, 2, 1]],
             sparse_size=(3, 3),
             sort_order='row',
             is_undirected=True,
@@ -210,7 +211,7 @@ class EdgeIndex(Tensor):
         assert not edge_index.is_undirected
 
         # Sparse-Dense Matrix Multiplication:
-        out = edge_index.flip(0) @ torch.randn(3, 16)
+        out = edge_index.flip(0) @ torch.randn(3, 16)
         assert out.size() == (3, 16)
     """
     # See "https://pytorch.org/docs/stable/notes/extending.html"
@@ -297,8 +298,7 @@ class EdgeIndex(Tensor):
                 indptr = None
             data = torch.stack([row, col], dim=0)
 
-        if (torch_geometric.typing.WITH_PT112
-                and data.layout == torch.sparse_csc):
+        if data.layout == torch.sparse_csc:
             row = data.row_indices()
             indptr = data.ccol_indices()
 
@@ -325,7 +325,7 @@ class EdgeIndex(Tensor):
             elif sparse_size[0] is None and sparse_size[1] is not None:
                 sparse_size = (sparse_size[1], sparse_size[1])
 
-        out = Tensor._make_wrapper_subclass(  # type: ignore
+        out = Tensor._make_wrapper_subclass(
             cls,
             size=data.size(),
             strides=data.stride(),
@@ -803,7 +803,7 @@ class EdgeIndex(Tensor):
 
         size = self.get_sparse_size()
         if value is not None and value.dim() > 1:
-            size = size + value.size()[1:]  # type: ignore
+            size = size + value.size()[1:]
 
         out = torch.full(size, fill_value, dtype=dtype, device=self.device)
         out[self._data[0], self._data[1]] = value if value is not None else 1
@@ -820,19 +820,28 @@ class EdgeIndex(Tensor):
                 :obj:`1.0`. (default: :obj:`None`)
         """
         value = self._get_value() if value is None else value
-        out = torch.sparse_coo_tensor(
+
+        if not torch_geometric.typing.WITH_PT21:
+            out = torch.sparse_coo_tensor(
+                indices=self._data,
+                values=value,
+                size=self.get_sparse_size(),
+                device=self.device,
+                requires_grad=value.requires_grad,
+            )
+            if self.is_sorted_by_row:
+                out = out._coalesced_(True)
+            return out
+
+        return torch.sparse_coo_tensor(
             indices=self._data,
             values=value,
             size=self.get_sparse_size(),
             device=self.device,
             requires_grad=value.requires_grad,
+            is_coalesced=True if self.is_sorted_by_row else None,
         )
 
-        if self.is_sorted_by_row:
-            out = out._coalesced_(True)
-
-        return out
-
     def to_sparse_csr(  # type: ignore
             self,
             value: Optional[Tensor] = None,
@@ -872,10 +881,6 @@ class EdgeIndex(Tensor):
                 If not specified, non-zero elements will be assigned a value of
                 :obj:`1.0`. (default: :obj:`None`)
         """
-        if not torch_geometric.typing.WITH_PT112:
-            raise NotImplementedError(
-                "'to_sparse_csc' not supported for PyTorch < 1.12")
-
         (colptr, row), perm = self.get_csc()
         if value is not None and perm is not None:
             value = value[perm]
@@ -912,7 +917,7 @@ class EdgeIndex(Tensor):
             return self.to_sparse_coo(value)
         if layout == torch.sparse_csr:
             return self.to_sparse_csr(value)
-        if torch_geometric.typing.WITH_PT112 and layout == torch.sparse_csc:
+        if layout == torch.sparse_csc:
             return self.to_sparse_csc(value)
 
         raise ValueError(f"Unexpected tensor layout (got '{layout}')")
@@ -1181,10 +1186,10 @@ class EdgeIndex(Tensor):
         return edge_index
 
     # Prevent auto-wrapping outputs back into the proper subclass type:
-    __torch_function__ = torch._C._disabled_torch_function_impl
+    __torch_function__ = torch._C._disabled_torch_function_impl  # type: ignore
 
     @classmethod
-    def __torch_dispatch__(
+    def __torch_dispatch__(  # type: ignore
         cls: Type,
         func: Callable[..., Any],
         types: Iterable[Type[Any]],
@@ -1237,6 +1242,14 @@ class EdgeIndex(Tensor):
         return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
                                                indent, force_newline=False)
 
+    def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
+        return self._data.tolist()
+
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self._data.numpy(force=force)
+
     # Helpers #################################################################
 
     def _shallow_copy(self) -> 'EdgeIndex':
@@ -1469,7 +1482,7 @@ def _slice(
     step: int = 1,
 ) -> Union[EdgeIndex, Tensor]:
 
-    if ((start is None or start <= 0)
+    if ((start is None or start == 0 or start <= -input.size(dim))
             and (end is None or end > input.size(dim)) and step == 1):
         return input._shallow_copy()  # No-op.
 
@@ -1928,7 +1941,7 @@ def _spmm(
     if transpose and not input.is_sorted_by_col:
         cls_name = input.__class__.__name__
         raise ValueError(f"'matmul(..., transpose=True)' requires "
-                         f"'{cls_name}' to be sorted by colums")
+                         f"'{cls_name}' to be sorted by columns")
 
     if (torch_geometric.typing.WITH_TORCH_SPARSE and not is_compiling()
             and other.is_cuda):  # pragma: no cover

torch_geometric/explain/algorithm/attention_explainer.py

@@ -1,13 +1,14 @@
 import logging
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union, overload
 
 import torch
 from torch import Tensor
 
-from torch_geometric.explain import Explanation
+from torch_geometric.explain import Explanation, HeteroExplanation
 from torch_geometric.explain.algorithm import ExplainerAlgorithm
 from torch_geometric.explain.config import ExplanationType, ModelTaskLevel
 from torch_geometric.nn.conv.message_passing import MessagePassing
+from torch_geometric.typing import EdgeType, NodeType
 
 
 class AttentionExplainer(ExplainerAlgorithm):
@@ -26,7 +27,9 @@ class AttentionExplainer(ExplainerAlgorithm):
     def __init__(self, reduce: str = 'max'):
         super().__init__()
         self.reduce = reduce
+        self.is_hetero = False
 
+    @overload
     def forward(
         self,
         model: torch.nn.Module,
@@ -37,65 +40,252 @@ class AttentionExplainer(ExplainerAlgorithm):
         index: Optional[Union[int, Tensor]] = None,
         **kwargs,
     ) -> Explanation:
-        if isinstance(x, dict) or isinstance(edge_index, dict):
-            raise ValueError(f"Heterogeneous graphs not yet supported in "
-                             f"'{self.__class__.__name__}'")
+        ...
 
-        hard_edge_mask = None
-        if self.model_config.task_level == ModelTaskLevel.node:
-            # We need to compute the hard edge mask to properly clean up edge
-            # attributions not involved during message passing:
-            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
-                                                     num_nodes=x.size(0))
+    @overload
+    def forward(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> HeteroExplanation:
+        ...
+
+    def forward(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> Union[Explanation, HeteroExplanation]:
+        """Generate explanations based on attention coefficients."""
+        self.is_hetero = isinstance(x, dict)
+
+        # Collect attention coefficients
+        alphas_dict = self._collect_attention_coefficients(
+            model, x, edge_index, **kwargs)
+
+        # Process attention coefficients
+        if self.is_hetero:
+            return self._create_hetero_explanation(model, alphas_dict,
+                                                   edge_index, index, x)
+        else:
+            return self._create_homo_explanation(model, alphas_dict,
+                                                 edge_index, index, x)
+
+    @overload
+    def _collect_attention_coefficients(
+        self,
+        model: torch.nn.Module,
+        x: Tensor,
+        edge_index: Tensor,
+        **kwargs,
+    ) -> List[Tensor]:
+        ...
+
+    @overload
+    def _collect_attention_coefficients(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        **kwargs,
+    ) -> Dict[EdgeType, List[Tensor]]:
+        ...
+
+    def _collect_attention_coefficients(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        **kwargs,
+    ) -> Union[List[Tensor], Dict[EdgeType, List[Tensor]]]:
+        """Collect attention coefficients from model layers."""
+        if self.is_hetero:
+            # For heterogeneous graphs, store alphas by edge type
+            alphas_dict: Dict[EdgeType, List[Tensor]] = {}
+
+            # Get list of edge types
+            edge_types = list(edge_index.keys())
+
+            # Hook function to capture attention coefficients by edge type
+            def hook(module, msg_kwargs, out):
+                # Find edge type from the module's full name
+                module_name = getattr(module, '_name', None)
+                if module_name is None:
+                    return
 
-        alphas: List[Tensor] = []
+                edge_type = None
+                for edge_tuple in edge_types:
+                    src_type, edge_name, dst_type = edge_tuple
+                    # Check if all components appear in the module name in
+                    # order
+                    try:
+                        src_idx = module_name.index(src_type)
+                        edge_idx = module_name.index(edge_name, src_idx)
+                        dst_idx = module_name.index(dst_type, edge_idx)
+                        if src_idx < edge_idx < dst_idx:
+                            edge_type = edge_tuple
+                            break
+                    except ValueError:  # Component not found
+                        continue
+
+                if edge_type is None:
+                    return
+
+                if edge_type not in alphas_dict:
+                    alphas_dict[edge_type] = []
+
+                # Extract alpha from message kwargs or module
+                if 'alpha' in msg_kwargs[0]:
+                    alphas_dict[edge_type].append(
+                        msg_kwargs[0]['alpha'].detach())
+                elif getattr(module, '_alpha', None) is not None:
+                    alphas_dict[edge_type].append(module._alpha.detach())
+        else:
+            # For homogeneous graphs, store all alphas in a list
+            alphas: List[Tensor] = []
 
-        def hook(module, msg_kwargs, out):
-            if 'alpha' in msg_kwargs[0]:
-                alphas.append(msg_kwargs[0]['alpha'].detach())
-            elif getattr(module, '_alpha', None) is not None:
-                alphas.append(module._alpha.detach())
+            def hook(module, msg_kwargs, out):
+                if 'alpha' in msg_kwargs[0]:
+                    alphas.append(msg_kwargs[0]['alpha'].detach())
+                elif getattr(module, '_alpha', None) is not None:
+                    alphas.append(module._alpha.detach())
 
+        # Register hooks for all message passing modules
         hook_handles = []
-        for module in model.modules():  # Register message forward hooks:
-            if (isinstance(module, MessagePassing)
-                    and module.explain is not False):
+        for name, module in model.named_modules():
+            if isinstance(module,
+                          MessagePassing) and module.explain is not False:
+                # Store name for hetero graph lookup in the hook
+                if self.is_hetero:
+                    module._name = name
+
                 hook_handles.append(module.register_message_forward_hook(hook))
 
+        # Forward pass to collect attention coefficients.
         model(x, edge_index, **kwargs)
 
-        for handle in hook_handles:  # Remove hooks:
+        # Remove hooks
+        for handle in hook_handles:
             handle.remove()
 
-        if len(alphas) == 0:
-            raise ValueError("Could not collect any attention coefficients. "
-                             "Please ensure that your model is using "
-                             "attention-based GNN layers.")
+        # Check if we collected any attention coefficients.
+        if self.is_hetero:
+            if not alphas_dict:
+                raise ValueError(
+                    "Could not collect any attention coefficients. "
+                    "Please ensure that your model is using "
+                    "attention-based GNN layers.")
+            return alphas_dict
+        else:
+            if not alphas:
+                raise ValueError(
+                    "Could not collect any attention coefficients. "
+                    "Please ensure that your model is using "
+                    "attention-based GNN layers.")
+            return alphas
 
+    def _process_attention_coefficients(
+        self,
+        alphas: List[Tensor],
+        edge_index_size: int,
+    ) -> Tensor:
+        """Process collected attention coefficients into a single mask."""
         for i, alpha in enumerate(alphas):
-            alpha = alpha[:edge_index.size(1)]  # Respect potential self-loops.
+            # Ensure alpha doesn't exceed edge_index size
+            alpha = alpha[:edge_index_size]
+
+            # Reduce multi-head attention
             if alpha.dim() == 2:
                 alpha = getattr(torch, self.reduce)(alpha, dim=-1)
-                if isinstance(alpha, tuple):  # Respect `torch.max`:
+                if isinstance(alpha, tuple):  # Handle torch.max output
                     alpha = alpha[0]
             elif alpha.dim() > 2:
-                raise ValueError(f"Can not reduce attention coefficients of "
+                raise ValueError(f"Cannot reduce attention coefficients of "
                                  f"shape {list(alpha.size())}")
             alphas[i] = alpha
 
+        # Combine attention coefficients across layers
         if len(alphas) > 1:
             alpha = torch.stack(alphas, dim=-1)
             alpha = getattr(torch, self.reduce)(alpha, dim=-1)
-            if isinstance(alpha, tuple):  # Respect `torch.max`:
+            if isinstance(alpha, tuple):  # Handle torch.max output
                 alpha = alpha[0]
         else:
             alpha = alphas[0]
 
+        return alpha
+
+    def _create_homo_explanation(
+        self,
+        model: torch.nn.Module,
+        alphas: List[Tensor],
+        edge_index: Tensor,
+        index: Optional[Union[int, Tensor]],
+        x: Tensor,
+    ) -> Explanation:
+        """Create explanation for homogeneous graph."""
+        # Get hard edge mask for node-level tasks
+        hard_edge_mask = None
+        if self.model_config.task_level == ModelTaskLevel.node:
+            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
+                                                     num_nodes=x.size(0))
+
+        # Process attention coefficients
+        alpha = self._process_attention_coefficients(alphas,
+                                                     edge_index.size(1))
+
+        # Post-process mask with hard edge mask if needed
         alpha = self._post_process_mask(alpha, hard_edge_mask,
                                         apply_sigmoid=False)
 
         return Explanation(edge_mask=alpha)
 
+    def _create_hetero_explanation(
+        self,
+        model: torch.nn.Module,
+        alphas_dict: Dict[EdgeType, List[Tensor]],
+        edge_index: Dict[EdgeType, Tensor],
+        index: Optional[Union[int, Tensor]],
+        x: Dict[NodeType, Tensor],
+    ) -> HeteroExplanation:
+        """Create explanation for heterogeneous graph."""
+        edge_masks_dict = {}
+
+        # Process each edge type separately
+        for edge_type, alphas in alphas_dict.items():
+            if not alphas:
+                continue
+
+            # Get hard edge mask for node-level tasks
+            hard_edge_mask = None
+            if self.model_config.task_level == ModelTaskLevel.node:
+                src_type, _, dst_type = edge_type
+                _, hard_edge_mask = self._get_hard_masks(
+                    model, index, edge_index[edge_type],
+                    num_nodes=max(x[src_type].size(0), x[dst_type].size(0)))
+
+            # Process attention coefficients for this edge type
+            alpha = self._process_attention_coefficients(
+                alphas, edge_index[edge_type].size(1))
+
+            # Apply hard mask if available
+            edge_masks_dict[edge_type] = self._post_process_mask(
+                alpha, hard_edge_mask, apply_sigmoid=False)
+
+        # Create heterogeneous explanation
+        explanation = HeteroExplanation()
+        explanation.set_value_dict('edge_mask', edge_masks_dict)
+        return explanation
+
     def supports(self) -> bool:
         explanation_type = self.explainer_config.explanation_type
         if explanation_type != ExplanationType.model:

torch_geometric/explain/algorithm/base.py

@@ -166,7 +166,7 @@ class ExplainerAlgorithm(torch.nn.Module):
         elif self.model_config.return_type == ModelReturnType.probs:
             loss_fn = F.binary_cross_entropy
         else:
-            assert False
+            raise AssertionError()
 
         return loss_fn(y_hat.view_as(y), y.float())
 
@@ -183,7 +183,7 @@ class ExplainerAlgorithm(torch.nn.Module):
         elif self.model_config.return_type == ModelReturnType.log_probs:
             loss_fn = F.nll_loss
         else:
-            assert False
+            raise AssertionError()
 
         return loss_fn(y_hat, y)
 

torch_geometric/explain/algorithm/captum.py

@@ -190,7 +190,7 @@ def to_captum_input(
 
     Args:
         x (torch.Tensor or Dict[NodeType, torch.Tensor]): The node features.
-            For heterogeneous graphs this is a dictionary holding node featues
+            For heterogeneous graphs this is a dictionary holding node features
             for each node type.
         edge_index(torch.Tensor or Dict[EdgeType, torch.Tensor]): The edge
             indices. For heterogeneous graphs this is a dictionary holding the

torch_geometric/explain/algorithm/captum_explainer.py

@@ -73,7 +73,8 @@ class CaptumExplainer(ExplainerAlgorithm):
                              f"{self.attribution_method_class.__name__}")
 
         if kwargs.get('internal_batch_size', 1) != 1:
-            warnings.warn("Overriding 'internal_batch_size' to 1")
+            warnings.warn("Overriding 'internal_batch_size' to 1",
+                          stacklevel=2)
 
         if 'internal_batch_size' in self._get_attribute_parameters():
             kwargs['internal_batch_size'] = 1