class_docstr: The Graph Neural Network Force Field (GNNFF) from the
  `"Accurate and scalable graph neural network force field and molecular
  dynamics with direct force architecture"
  <https://www.nature.com/articles/s41524-021-00543-3>`_ paper.
  :class:`GNNFF` directly predicts atomic forces from automatically
  extracted features of the local atomic environment that are
  translationally-invariant, but rotationally-covariant to the coordinate of
  the atoms.
  
  Args:
      hidden_node_channels (int): Hidden node embedding size.
      hidden_edge_channels (int): Hidden edge embedding size.
      num_layers (int): Number of message passing blocks.
      cutoff (float, optional): Cutoff distance for interatomic
          interactions. (default: :obj:`5.0`)
      max_num_neighbors (int, optional): The maximum number of neighbors to
          collect for each node within the :attr:`cutoff` distance.
          (default: :obj:`32`)
code: def __init__(
      self,
      hidden_node_channels: int,
      hidden_edge_channels: int,
      num_layers: int,
      cutoff: float = 5.0,
      max_num_neighbors: int = 32,
  ):
      super().__init__()
  
      self.cutoff = cutoff
      self.max_num_neighbors = max_num_neighbors
  
      self.node_emb = Sequential(
          Embedding(95, hidden_node_channels),
          ShiftedSoftplus(),
          Linear(hidden_node_channels, hidden_node_channels),
          ShiftedSoftplus(),
          Linear(hidden_node_channels, hidden_node_channels),
      )
      self.edge_emb = GaussianFilter(0.0, 5.0, hidden_edge_channels)
  
      self.node_blocks = ModuleList(
          [
              NodeBlock(hidden_node_channels, hidden_edge_channels)
              for _ in range(num_layers)
          ]
      )
      self.edge_blocks = ModuleList(
          [
              EdgeBlock(hidden_node_channels, hidden_edge_channels)
              for _ in range(num_layers)
          ]
      )
  
      self.force_predictor = Sequential(
          Linear(hidden_edge_channels, hidden_edge_channels),
          ShiftedSoftplus(),
          Linear(hidden_edge_channels, hidden_edge_channels),
          ShiftedSoftplus(),
          Linear(hidden_edge_channels, 1),
      )
  
imports: import torch
  from torch import Tensor
  from torch.nn import BatchNorm1d, Embedding, Linear, ModuleList, Sequential
  from torch_geometric.nn import radius_graph
  from torch_geometric.nn.inits import reset
  from torch_geometric.nn.models.dimenet import triplets
  from torch_geometric.nn.models.schnet import ShiftedSoftplus
  from torch_geometric.typing import OptTensor
  from torch_geometric.utils import scatter
label_desc: Use image node_id 1 to create a new GNNFF object from inherited base classes: torch.nn with example: obj = GNNFF(hidden_node_channels, hidden_edge_channels, num_layers, cutoff, max_num_neighbors)