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 _get_updated_memory(self, n_id: Tensor) -> Tuple[Tensor, Tensor]:
      self._assoc[n_id] = torch.arange(n_id.size(0), device=n_id.device)
  
      # Compute messages (src -> dst).
      msg_s, t_s, src_s, dst_s = self._compute_msg(
          n_id, self.msg_s_store, self.msg_s_module
      )
  
      # Compute messages (dst -> src).
      msg_d, t_d, src_d, dst_d = self._compute_msg(
          n_id, self.msg_d_store, self.msg_d_module
      )
  
      # Aggregate messages.
      idx = torch.cat([src_s, src_d], dim=0)
      msg = torch.cat([msg_s, msg_d], dim=0)
      t = torch.cat([t_s, t_d], dim=0)
      aggr = self.aggr_module(msg, self._assoc[idx], t, n_id.size(0))
  
      # Get local copy of updated memory.
      memory = self.gru(aggr, self.memory[n_id])
  
      # Get local copy of updated `last_update`.
      dim_size = self.last_update.size(0)
      last_update = scatter(t, idx, 0, dim_size, reduce="max")[n_id]
  
      return memory, last_update
  
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 10 for calling the TGNMemory obj's underlying member method code with example usage: obj._get_updated_memory(n_id) and returns: memory,  last_update