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

torch_geometric/nn/norm/layer_norm.py

@@ -30,11 +30,13 @@ class LayerNorm(torch.nn.Module):
         affine (bool, optional): If set to :obj:`True`, this module has
             learnable affine parameters :math:`\gamma` and :math:`\beta`.
             (default: :obj:`True`)
-        mode (str, optinal): The normalization mode to use for layer
+        mode (str, optional): The normalization mode to use for layer
             normalization (:obj:`"graph"` or :obj:`"node"`). If :obj:`"graph"`
             is used, each graph will be considered as an element to be
             normalized. If `"node"` is used, each node will be considered as
             an element to be normalized. (default: :obj:`"graph"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -42,6 +44,7 @@ class LayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'graph',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -51,8 +54,8 @@ class LayerNorm(torch.nn.Module):
         self.mode = mode
 
         if affine:
-            self.weight = Parameter(torch.empty(in_channels))
-            self.bias = Parameter(torch.empty(in_channels))
+            self.weight = Parameter(torch.empty(in_channels, device=device))
+            self.bias = Parameter(torch.empty(in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
@@ -108,7 +111,7 @@ class LayerNorm(torch.nn.Module):
             return F.layer_norm(x, (self.in_channels, ), self.weight,
                                 self.bias, self.eps)
 
-        raise ValueError(f"Unknow normalization mode: {self.mode}")
+        raise ValueError(f"Unknownn normalization mode: {self.mode}")
 
     def __repr__(self):
         return (f'{self.__class__.__name__}({self.in_channels}, '
@@ -130,10 +133,12 @@ class HeteroLayerNorm(torch.nn.Module):
         affine (bool, optional): If set to :obj:`True`, this module has
             learnable affine parameters :math:`\gamma` and :math:`\beta`.
             (default: :obj:`True`)
-        mode (str, optinal): The normalization mode to use for layer
+        mode (str, optional): The normalization mode to use for layer
             normalization (:obj:`"node"`). If `"node"` is used, each node will
             be considered as an element to be normalized.
             (default: :obj:`"node"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -142,6 +147,7 @@ class HeteroLayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'node',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
         assert mode == 'node'
@@ -152,8 +158,10 @@ class HeteroLayerNorm(torch.nn.Module):
         self.affine = affine
 
         if affine:
-            self.weight = Parameter(torch.empty(num_types, in_channels))
-            self.bias = Parameter(torch.empty(num_types, in_channels))
+            self.weight = Parameter(
+                torch.empty(num_types, in_channels, device=device))
+            self.bias = Parameter(
+                torch.empty(num_types, in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)

torch_geometric/nn/norm/msg_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -19,10 +21,14 @@ class MessageNorm(torch.nn.Module):
         learn_scale (bool, optional): If set to :obj:`True`, will learn the
             scaling factor :math:`s` of message normalization.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, learn_scale: bool = False):
+    def __init__(self, learn_scale: bool = False,
+                 device: Optional[torch.device] = None):
         super().__init__()
-        self.scale = Parameter(torch.empty(1), requires_grad=learn_scale)
+        self.scale = Parameter(torch.empty(1, device=device),
+                               requires_grad=learn_scale)
         self.reset_parameters()
 
     def reset_parameters(self):

torch_geometric/nn/pool/__init__.py

@@ -7,18 +7,19 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric.typing import OptTensor, torch_cluster
 
-from .asap import ASAPooling
 from .avg_pool import avg_pool, avg_pool_neighbor_x, avg_pool_x
-from .edge_pool import EdgePooling
 from .glob import global_add_pool, global_max_pool, global_mean_pool
 from .knn import (KNNIndex, L2KNNIndex, MIPSKNNIndex, ApproxL2KNNIndex,
                   ApproxMIPSKNNIndex)
 from .graclus import graclus
 from .max_pool import max_pool, max_pool_neighbor_x, max_pool_x
-from .mem_pool import MemPooling
-from .pan_pool import PANPooling
-from .sag_pool import SAGPooling
 from .topk_pool import TopKPooling
+from .sag_pool import SAGPooling
+from .edge_pool import EdgePooling
+from .cluster_pool import ClusterPooling
+from .asap import ASAPooling
+from .pan_pool import PANPooling
+from .mem_pool import MemPooling
 from .voxel_grid import voxel_grid
 from .approx_knn import approx_knn, approx_knn_graph
 
@@ -162,8 +163,10 @@ def knn_graph(
     :rtype: :class:`torch.Tensor`
     """
     if batch is not None and x.device != batch.device:
-        warnings.warn("Input tensor 'x' and 'batch' are on different devices "
-                      "in 'knn_graph'. Performing blocking device transfer")
+        warnings.warn(
+            "Input tensor 'x' and 'batch' are on different devices "
+            "in 'knn_graph'. Performing blocking device transfer",
+            stacklevel=2)
         batch = batch.to(x.device)
 
     if not torch_geometric.typing.WITH_TORCH_CLUSTER_BATCH_SIZE:
@@ -284,8 +287,10 @@ def radius_graph(
         inputs to GPU before proceeding.
     """
     if batch is not None and x.device != batch.device:
-        warnings.warn("Input tensor 'x' and 'batch' are on different devices "
-                      "in 'radius_graph'. Performing blocking device transfer")
+        warnings.warn(
+            "Input tensor 'x' and 'batch' are on different devices "
+            "in 'radius_graph'. Performing blocking device transfer",
+            stacklevel=2)
         batch = batch.to(x.device)
 
     if not torch_geometric.typing.WITH_TORCH_CLUSTER_BATCH_SIZE:
@@ -344,6 +349,7 @@ __all__ = [
     'TopKPooling',
     'SAGPooling',
     'EdgePooling',
+    'ClusterPooling',
     'ASAPooling',
     'PANPooling',
     'MemPooling',

torch_geometric/nn/pool/cluster_pool.py

@@ -0,0 +1,144 @@
+from typing import NamedTuple, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.utils import (
+    dense_to_sparse,
+    one_hot,
+    to_dense_adj,
+    to_scipy_sparse_matrix,
+)
+
+
+class UnpoolInfo(NamedTuple):
+    edge_index: Tensor
+    cluster: Tensor
+    batch: Tensor
+
+
+class ClusterPooling(torch.nn.Module):
+    r"""The cluster pooling operator from the `"Edge-Based Graph Component
+    Pooling" <https://arxiv.org/abs/2409.11856>`_ paper.
+    :class:`ClusterPooling` computes a score for each edge.
+    Based on the selected edges, graph clusters are calculated and compressed
+    to one node using the injective :obj:`"sum"` aggregation function.
+    Edges are remapped based on the nodes created by each cluster and the
+    original edges.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        edge_score_method (str, optional): The function to apply
+            to compute the edge score from raw edge scores (:obj:`"tanh"`,
+            :obj:`"sigmoid"`, :obj:`"log_softmax"`). (default: :obj:`"tanh"`)
+        dropout (float, optional): The probability with
+            which to drop edge scores during training. (default: :obj:`0.0`)
+        threshold (float, optional): The threshold of edge scores. If set to
+            :obj:`None`, will be automatically inferred depending on
+            :obj:`edge_score_method`. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        edge_score_method: str = 'tanh',
+        dropout: float = 0.0,
+        threshold: Optional[float] = None,
+    ):
+        super().__init__()
+        assert edge_score_method in ['tanh', 'sigmoid', 'log_softmax']
+
+        if threshold is None:
+            threshold = 0.5 if edge_score_method == 'sigmoid' else 0.0
+
+        self.in_channels = in_channels
+        self.edge_score_method = edge_score_method
+        self.dropout = dropout
+        self.threshold = threshold
+
+        self.lin = torch.nn.Linear(2 * in_channels, 1)
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.lin.reset_parameters()
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): The node features.
+            edge_index (torch.Tensor): The edge indices.
+            batch (torch.Tensor): Batch vector
+                :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
+                each node to a specific example.
+
+        Return types:
+            * **x** *(torch.Tensor)* - The pooled node features.
+            * **edge_index** *(torch.Tensor)* - The coarsened edge indices.
+            * **batch** *(torch.Tensor)* - The coarsened batch vector.
+            * **unpool_info** *(UnpoolInfo)* - Information that can be consumed
+              for unpooling.
+        """
+        mask = edge_index[0] != edge_index[1]
+        edge_index = edge_index[:, mask]
+
+        edge_attr = torch.cat(
+            [x[edge_index[0]], x[edge_index[1]]],
+            dim=-1,
+        )
+        edge_score = self.lin(edge_attr).view(-1)
+        edge_score = F.dropout(edge_score, p=self.dropout,
+                               training=self.training)
+
+        if self.edge_score_method == 'tanh':
+            edge_score = edge_score.tanh()
+        elif self.edge_score_method == 'sigmoid':
+            edge_score = edge_score.sigmoid()
+        else:
+            assert self.edge_score_method == 'log_softmax'
+            edge_score = F.log_softmax(edge_score, dim=0)
+
+        return self._merge_edges(x, edge_index, batch, edge_score)
+
+    def _merge_edges(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_score: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]:
+
+        from scipy.sparse.csgraph import connected_components
+
+        edge_contract = edge_index[:, edge_score > self.threshold]
+
+        adj = to_scipy_sparse_matrix(edge_contract, num_nodes=x.size(0))
+        _, cluster_np = connected_components(adj, directed=True,
+                                             connection="weak")
+
+        cluster = torch.tensor(cluster_np, dtype=torch.long, device=x.device)
+        C = one_hot(cluster)
+        A = to_dense_adj(edge_index, max_num_nodes=x.size(0)).squeeze(0)
+        S = to_dense_adj(edge_index, edge_attr=edge_score,
+                         max_num_nodes=x.size(0)).squeeze(0)
+
+        A_contract = to_dense_adj(edge_contract,
+                                  max_num_nodes=x.size(0)).squeeze(0)
+        nodes_single = ((A_contract.sum(dim=-1) +
+                         A_contract.sum(dim=-2)) == 0).nonzero()
+        S[nodes_single, nodes_single] = 1.0
+
+        x_out = (S @ C).t() @ x
+        edge_index_out, _ = dense_to_sparse((C.T @ A @ C).fill_diagonal_(0))
+        batch_out = batch.new_empty(x_out.size(0)).scatter_(0, cluster, batch)
+        unpool_info = UnpoolInfo(edge_index, cluster, batch)
+
+        return x_out, edge_index_out, batch_out, unpool_info
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.in_channels})'

torch_geometric/nn/pool/connect/base.py

@@ -4,7 +4,6 @@ from typing import Optional
 import torch
 from torch import Tensor
 
-import torch_geometric.typing
 from torch_geometric.nn.pool.select import SelectOutput
 
 
@@ -49,8 +48,7 @@ class ConnectOutput:
         self.batch = batch
 
 
-if torch_geometric.typing.WITH_PT113:
-    ConnectOutput = torch.jit.script(ConnectOutput)
+ConnectOutput = torch.jit.script(ConnectOutput)
 
 
 class Connect(torch.nn.Module):

torch_geometric/nn/pool/edge_pool.py

@@ -58,7 +58,7 @@ class EdgePooling(torch.nn.Module):
         self,
         in_channels: int,
         edge_score_method: Optional[Callable] = None,
-        dropout: Optional[float] = 0.0,
+        dropout: float = 0.0,
         add_to_edge_score: float = 0.5,
     ):
         super().__init__()

torch_geometric/nn/pool/knn.py

@@ -91,9 +91,10 @@ class KNNIndex:
                     if hasattr(self.index, 'reserveMemory'):
                         self.index.reserveMemory(self.reserve)
                     else:
-                        warnings.warn(f"'{self.index.__class__.__name__}' "
-                                      f"does not support pre-allocation of "
-                                      f"memory")
+                        warnings.warn(
+                            f"'{self.index.__class__.__name__}' "
+                            f"does not support pre-allocation of "
+                            f"memory", stacklevel=2)
 
             self.index.train(emb)
 
@@ -135,14 +136,16 @@ class KNNIndex:
         query_k = min(query_k, self.numel)
 
         if k > 2048:  # `faiss` supports up-to `k=2048`:
-            warnings.warn(f"Capping 'k' to faiss' upper limit of 2048 "
-                          f"(got {k}). This may cause some relevant items to "
-                          f"not be retrieved.")
+            warnings.warn(
+                f"Capping 'k' to faiss' upper limit of 2048 "
+                f"(got {k}). This may cause some relevant items to "
+                f"not be retrieved.", stacklevel=2)
         elif query_k > 2048:
-            warnings.warn(f"Capping 'k' to faiss' upper limit of 2048 "
-                          f"(got {k} which got extended to {query_k} due to "
-                          f"the exclusion of existing links). This may cause "
-                          f"some relevant items to not be retrieved.")
+            warnings.warn(
+                f"Capping 'k' to faiss' upper limit of 2048 "
+                f"(got {k} which got extended to {query_k} due to "
+                f"the exclusion of existing links). This may cause "
+                f"some relevant items to not be retrieved.", stacklevel=2)
             query_k = 2048
 
         score, index = self.index.search(emb.detach(), query_k)

torch_geometric/nn/pool/select/base.py

@@ -4,8 +4,6 @@ from typing import Optional
 import torch
 from torch import Tensor
 
-import torch_geometric.typing
-
 
 @dataclass(init=False)
 class SelectOutput:
@@ -64,8 +62,7 @@ class SelectOutput:
         self.weight = weight
 
 
-if torch_geometric.typing.WITH_PT113:
-    SelectOutput = torch.jit.script(SelectOutput)
+SelectOutput = torch.jit.script(SelectOutput)
 
 
 class Select(torch.nn.Module):

torch_geometric/nn/summary.py

@@ -141,7 +141,7 @@ def get_shape(inputs: Any) -> str:
 def postprocess(info_list: List[dict]) -> List[dict]:
     for idx, info in enumerate(info_list):
         depth = info['depth']
-        if idx > 0:  # root module (0) is exclued
+        if idx > 0:  # root module (0) is excluded
             if depth == 1:
                 prefix = '├─'
             else:

torch_geometric/nn/to_hetero_module.py

@@ -108,9 +108,10 @@ class ToHeteroMessagePassing(torch.nn.Module):
 
         if (not hasattr(module, 'reset_parameters')
                 and sum([p.numel() for p in module.parameters()]) > 0):
-            warnings.warn(f"'{module}' will be duplicated, but its parameters "
-                          f"cannot be reset. To suppress this warning, add a "
-                          f"'reset_parameters()' method to '{module}'")
+            warnings.warn(
+                f"'{module}' will be duplicated, but its parameters "
+                f"cannot be reset. To suppress this warning, add a "
+                f"'reset_parameters()' method to '{module}'", stacklevel=2)
 
         convs = {edge_type: copy.deepcopy(module) for edge_type in edge_types}
         self.hetero_module = HeteroConv(convs, aggr)

torch_geometric/nn/to_hetero_transformer.py

@@ -157,7 +157,7 @@ class ToHeteroTransformer(Transformer):
                 f"There exist node types ({unused_node_types}) whose "
                 f"representations do not get updated during message passing "
                 f"as they do not occur as destination type in any edge type. "
-                f"This may lead to unexpected behavior.")
+                f"This may lead to unexpected behavior.", stacklevel=2)
 
         names = self.metadata[0] + [rel for _, rel, _ in self.metadata[1]]
         for name in names:
@@ -166,7 +166,7 @@ class ToHeteroTransformer(Transformer):
                     f"The type '{name}' contains invalid characters which "
                     f"may lead to unexpected behavior. To avoid any issues, "
                     f"ensure that your types only contain letters, numbers "
-                    f"and underscores.")
+                    f"and underscores.", stacklevel=2)
 
     def placeholder(self, node: Node, target: Any, name: str):
         # Adds a `get` call to the input dictionary for every node-type or
@@ -379,7 +379,7 @@ class ToHeteroTransformer(Transformer):
                 warnings.warn(
                     f"'{target}' will be duplicated, but its parameters "
                     f"cannot be reset. To suppress this warning, add a "
-                    f"'reset_parameters()' method to '{target}'")
+                    f"'reset_parameters()' method to '{target}'", stacklevel=2)
 
         return module_dict
 

torch_geometric/nn/to_hetero_with_bases_transformer.py

@@ -165,7 +165,7 @@ class ToHeteroWithBasesTransformer(Transformer):
                 f"There exist node types ({unused_node_types}) whose "
                 f"representations do not get updated during message passing "
                 f"as they do not occur as destination type in any edge type. "
-                f"This may lead to unexpected behavior.")
+                f"This may lead to unexpected behavior.", stacklevel=2)
 
         names = self.metadata[0] + [rel for _, rel, _ in self.metadata[1]]
         for name in names:
@@ -174,7 +174,7 @@ class ToHeteroWithBasesTransformer(Transformer):
                     f"The type '{name}' contains invalid characters which "
                     f"may lead to unexpected behavior. To avoid any issues, "
                     f"ensure that your types only contain letters, numbers "
-                    f"and underscores.")
+                    f"and underscores.", stacklevel=2)
 
     def transform(self) -> GraphModule:
         self._node_offset_dict_initialized = False
@@ -361,7 +361,7 @@ class HeteroBasisConv(torch.nn.Module):
                 warnings.warn(
                     f"'{conv}' will be duplicated, but its parameters cannot "
                     f"be reset. To suppress this warning, add a "
-                    f"'reset_parameters()' method to '{conv}'")
+                    f"'reset_parameters()' method to '{conv}'", stacklevel=2)
             torch.nn.init.xavier_uniform_(conv.edge_type_weight)
 
     def forward(self, edge_type: Tensor, *args, **kwargs) -> Tensor:
@@ -380,7 +380,7 @@ class HeteroBasisConv(torch.nn.Module):
 
 
 class LinearAlign(torch.nn.Module):
-    # Aligns representions to the same dimensionality. Note that this will
+    # Aligns representations to the same dimensionality. Note that this will
     # create lazy modules, and as such requires a forward pass in order to
     # initialize parameters.
     def __init__(self, keys: List[Union[NodeType, EdgeType]],
@@ -468,7 +468,7 @@ def get_edge_type(
 ###############################################################################
 
 # These methods are used to group the individual type-wise components into a
-# unfied single representation.
+# unified single representation.
 
 
 def group_node_placeholder(input_dict: Dict[NodeType, Tensor],

torch_geometric/profile/__init__.py

@@ -20,6 +20,7 @@ from .utils import (
     get_gpu_memory_from_nvidia_smi,
     get_model_size,
 )
+from .nvtx import nvtxit
 
 __all__ = [
     'profileit',
@@ -38,6 +39,7 @@ __all__ = [
     'get_gpu_memory_from_nvidia_smi',
     'get_gpu_memory_from_ipex',
     'benchmark',
+    'nvtxit',
 ]
 
 classes = __all__

torch_geometric/profile/nvtx.py

@@ -0,0 +1,66 @@
+from functools import wraps
+from typing import Optional
+
+import torch
+
+CUDA_PROFILE_STARTED = False
+
+
+def begin_cuda_profile():
+    global CUDA_PROFILE_STARTED
+    prev_state = CUDA_PROFILE_STARTED
+    if prev_state is False:
+        CUDA_PROFILE_STARTED = True
+        torch.cuda.cudart().cudaProfilerStart()
+    return prev_state
+
+
+def end_cuda_profile(prev_state: bool):
+    global CUDA_PROFILE_STARTED
+    CUDA_PROFILE_STARTED = prev_state
+    if prev_state is False:
+        torch.cuda.cudart().cudaProfilerStop()
+
+
+def nvtxit(name: Optional[str] = None, n_warmups: int = 0,
+           n_iters: Optional[int] = None):
+    """Enables NVTX profiling for a function.
+
+    Args:
+        name (Optional[str], optional): Name to give the reference frame for
+            the function being wrapped. Defaults to the name of the
+            function in code.
+        n_warmups (int, optional): Number of iters to call that function
+            before starting. Defaults to 0.
+        n_iters (Optional[int], optional): Number of iters of that function to
+            record. Defaults to all of them.
+    """
+    def nvtx(func):
+
+        nonlocal name
+        iters_so_far = 0
+        if name is None:
+            name = func.__name__
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            nonlocal iters_so_far
+            if not torch.cuda.is_available():
+                return func(*args, **kwargs)
+            elif iters_so_far < n_warmups:
+                iters_so_far += 1
+                return func(*args, **kwargs)
+            elif n_iters is None or iters_so_far < n_iters + n_warmups:
+                prev_state = begin_cuda_profile()
+                torch.cuda.nvtx.range_push(f"{name}_{iters_so_far}")
+                result = func(*args, **kwargs)
+                torch.cuda.nvtx.range_pop()
+                end_cuda_profile(prev_state)
+                iters_so_far += 1
+                return result
+            else:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return nvtx

torch_geometric/profile/profiler.py

@@ -5,6 +5,8 @@ from typing import Any, List, NamedTuple, Optional, Tuple
 import torch
 import torch.profiler as torch_profiler
 
+import torch_geometric.typing
+
 # predefined namedtuple for variable setting (global template)
 Trace = namedtuple('Trace', ['path', 'leaf', 'module'])
 
@@ -325,6 +327,8 @@ def _flatten_tree(t, depth=0):
 
 
 def _build_measure_tuple(events: List, occurrences: List) -> NamedTuple:
+    device_str = 'device' if torch_geometric.typing.WITH_PT24 else 'cuda'
+
     # memory profiling supported in torch >= 1.6
     self_cpu_memory = None
     has_self_cpu_memory = any(
@@ -339,29 +343,34 @@ def _build_measure_tuple(events: List, occurrences: List) -> NamedTuple:
             [getattr(e, "cpu_memory_usage", 0) or 0 for e in events])
     self_cuda_memory = None
     has_self_cuda_memory = any(
-        hasattr(e, "self_cuda_memory_usage") for e in events)
+        hasattr(e, f"self_{device_str}_memory_usage") for e in events)
     if has_self_cuda_memory:
-        self_cuda_memory = sum(
-            [getattr(e, "self_cuda_memory_usage", 0) or 0 for e in events])
+        self_cuda_memory = sum([
+            getattr(e, f"self_{device_str}_memory_usage", 0) or 0
+            for e in events
+        ])
     cuda_memory = None
-    has_cuda_memory = any(hasattr(e, "cuda_memory_usage") for e in events)
+    has_cuda_memory = any(
+        hasattr(e, f"{device_str}_memory_usage") for e in events)
     if has_cuda_memory:
         cuda_memory = sum(
-            [getattr(e, "cuda_memory_usage", 0) or 0 for e in events])
+            [getattr(e, f"{device_str}_memory_usage", 0) or 0 for e in events])
 
     # self CUDA time supported in torch >= 1.7
     self_cuda_total = None
     has_self_cuda_time = any(
-        hasattr(e, "self_cuda_time_total") for e in events)
+        hasattr(e, f"self_{device_str}_time_total") for e in events)
     if has_self_cuda_time:
-        self_cuda_total = sum(
-            [getattr(e, "self_cuda_time_total", 0) or 0 for e in events])
+        self_cuda_total = sum([
+            getattr(e, f"self_{device_str}_time_total", 0) or 0 for e in events
+        ])
 
     return Measure(
         self_cpu_total=sum([e.self_cpu_time_total or 0 for e in events]),
         cpu_total=sum([e.cpu_time_total or 0 for e in events]),
         self_cuda_total=self_cuda_total,
-        cuda_total=sum([e.cuda_time_total or 0 for e in events]),
+        cuda_total=sum(
+            [getattr(e, f"{device_str}_time_total") or 0 for e in events]),
         self_cpu_memory=self_cpu_memory,
         cpu_memory=cpu_memory,
         self_cuda_memory=self_cuda_memory,