class_docstr: The Temporal Graph Network (TGN) memory model from the
  `"Temporal Graph Networks for Deep Learning on Dynamic Graphs"
  <https://arxiv.org/abs/2006.10637>`_ paper.
  
  .. note::
  
      For an example of using TGN, see `examples/tgn.py
      <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
      tgn.py>`_.
  
  Args:
      num_nodes (int): The number of nodes to save memories for.
      raw_msg_dim (int): The raw message dimensionality.
      memory_dim (int): The hidden memory dimensionality.
      time_dim (int): The time encoding dimensionality.
      message_module (torch.nn.Module): The message function which
          combines source and destination node memory embeddings, the raw
          message and the time encoding.
      aggregator_module (torch.nn.Module): The message aggregator function
          which aggregates messages to the same destination into a single
          representation.
code: def __init__(
      self,
      num_nodes: int,
      raw_msg_dim: int,
      memory_dim: int,
      time_dim: int,
      message_module: Callable,
      aggregator_module: Callable,
  ):
      super().__init__()
  
      self.num_nodes = num_nodes
      self.raw_msg_dim = raw_msg_dim
      self.memory_dim = memory_dim
      self.time_dim = time_dim
  
      self.msg_s_module = message_module
      self.msg_d_module = copy.deepcopy(message_module)
      self.aggr_module = aggregator_module
      self.time_enc = TimeEncoder(time_dim)
      self.gru = GRUCell(message_module.out_channels, memory_dim)
  
      self.register_buffer("memory", torch.empty(num_nodes, memory_dim))
      last_update = torch.empty(self.num_nodes, dtype=torch.long)
      self.register_buffer("last_update", last_update)
      self.register_buffer(
          "_assoc", torch.empty(num_nodes, dtype=torch.long)
      )
  
      self.msg_s_store = {}
      self.msg_d_store = {}
  
      self.reset_parameters()
  
imports: import copy
  from typing import Callable, Dict, Tuple
  import torch
  from torch import Tensor
  from torch.nn import GRUCell, Linear
  from torch_geometric.nn.inits import zeros
  from torch_geometric.utils import scatter
  from torch_geometric.utils._scatter import scatter_argmax
label_desc: Use image node_id 1 to create a new TGNMemory object from inherited base classes: torch.nn with example: obj = TGNMemory(num_nodes, raw_msg_dim, memory_dim, time_dim, message_module, aggregator_module)