class_docstr: MatErials Graph Network
  
  A model for predicting crystal and molecular properties using GraphNetworks.
  
  Example
  -------
  >>> import numpy as np
  >>> from torch_geometric.data import Batch
  >>> from deepchem.feat import GraphData
  >>> n_nodes, n_node_features = 5, 10
  >>> n_edges, n_edge_attrs = 5, 2
  >>> n_global_features = 4
  >>> node_features = np.random.randn(n_nodes, n_node_features)
  >>> edge_attrs = np.random.randn(n_edges, n_edge_attrs)
  >>> edge_index = np.array([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]], dtype=np.int64)
  >>> global_features = np.random.randn(1, n_global_features)
  >>> graph = GraphData(node_features, edge_index, edge_attrs, global_features=global_features)
  >>> batch = Batch()
  >>> batch = batch.from_data_list([graph.to_pyg_graph()])
  >>> model = MEGNet(n_node_features=n_node_features, n_edge_features=n_edge_attrs, n_global_features=n_global_features)
  >>> pred = model(batch)
  
  Note
  ----
  This class requires torch-geometric to be installed.
docstr: Parameters
  ----------
  pyg_batch: torch_geometric.data.Batch
      A pytorch-geometric batch of graphs where node attributes are stores
      as pyg_batch['x'], edge_index in pyg_batch['edge_index'], edge features
      in pyg_batch['edge_attr'], global features in pyg_batch['global_features']
  
  Returns
  -------
  torch.Tensor: Predictions for the graph
docstr_tok: ["Parameters","--","--","--","--","--","pyg_batch",":","torch_geometric.data.Batch","A","pytorch-geometric","batch","of","graphs","where","node","attributes","are","stores","as","pyg_batch","[","'","x","'","]",",","edge_index","in","pyg_batch","[","'edge_index","'","]",",","edge","features","in","pyg_batch","[","'edge_attr","'","]",",","global","features","in","pyg_batch","[","'global_features","'","]","Returns","--","--","--","-","torch.Tensor",":","Predictions","for","the","graph"]
code: def forward(self, pyg_batch):
      """
      Parameters
      ----------
      pyg_batch: torch_geometric.data.Batch
          A pytorch-geometric batch of graphs where node attributes are stores
          as pyg_batch['x'], edge_index in pyg_batch['edge_index'], edge features
          in pyg_batch['edge_attr'], global features in pyg_batch['global_features']
  
      Returns
      -------
      torch.Tensor: Predictions for the graph
      """
      node_features = pyg_batch["x"]
      edge_index, edge_features = (
          pyg_batch["edge_index"],
          pyg_batch["edge_attr"],
      )
      global_features = pyg_batch["global_features"]
      batch = pyg_batch["batch"]
  
      for i in range(self.n_blocks):
          (
              node_features,
              edge_features,
              global_features,
          ) = self.megnet_blocks[i](
              node_features,
              edge_index,
              edge_features,
              global_features,
              batch,
          )
  
      node_features = self.set2set_nodes(node_features, batch)
      edge_features = self.set2set_edges(
          edge_features, batch[edge_index[0]]
      )
      out = torch.cat(
          [node_features, edge_features, global_features], axis=1
      )
      out = self.out(self.dense(out))
  
      if self.mode == "classification":
          if self.n_tasks == 1:
              logits = out.view(-1, self.n_classes)
              softmax_dim = 1
          else:
              logits = out.view(-1, self.n_tasks, self.n_classes)
              softmax_dim = 2
          proba = F.softmax(logits, dim=softmax_dim)
          return proba, logits
      elif self.mode == "regression":
          return out
  
file_path: deepchem/deepchem/models/torch_models/megnet.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from deepchem.models.losses import Loss, L2Loss, SparseSoftmaxCrossEntropy
  from deepchem.models.torch_models.layers import GraphNetwork
  from deepchem.models.torch_models import TorchModel
label_desc: Use image node_id 2 for calling the MEGNet obj's underlying member method code with example usage: obj.forward(pyg_batch) and returns: proba,  logits, out