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 device(self) -> torch.device:
      return self.time_enc.lin.weight.device
  
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 2 for calling the TGNMemory obj's underlying member method code with example usage: obj.device() and returns: self