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 forward(
      self, z: Tensor, pos: Tensor, batch: OptTensor = None
  ) -> Tensor:
      """"""  # noqa: D419
      edge_index = radius_graph(
          pos,
          r=self.cutoff,
          batch=batch,
          max_num_neighbors=self.max_num_neighbors,
      )
  
      i, j, idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(
          edge_index, num_nodes=z.size(0)
      )
  
      # Calculate distances and unit vector:
      dist = (pos[i] - pos[j]).pow(2).sum(dim=-1).sqrt()
      unit_vec = (pos[i] - pos[j]) / dist.view(-1, 1)
  
      # Embedding blocks:
      node_emb = self.node_emb(z)
      edge_emb = self.edge_emb(dist)
  
      # Message passing blocks:
      for node_block, edge_block in zip(
          self.node_blocks, self.edge_blocks
      ):
          node_emb = node_block(node_emb, edge_emb, i)
          edge_emb = edge_block(
              node_emb,
              edge_emb,
              i,
              j,
              idx_i,
              idx_j,
              idx_k,
              idx_ji,
              idx_kj,
          )
  
      # Force prediction block:
      force = self.force_predictor(edge_emb) * unit_vec
  
      return scatter(force, i, dim=0, reduce="sum")
  
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 3 for calling the GNNFF obj's underlying member method code with example usage: obj.forward(z, pos, batch) and returns: scatter