class_docstr: MatErials Graph Network for Molecules and Crystals
  
  MatErials Graph Network [1]_ are Graph Networks [2]_ which are used for property prediction
  in molecules and crystals. The model implements multiple layers of Graph Network as
  MEGNetBlocks and then combines the node properties and edge properties of all nodes
  and edges via a Set2Set layer. The combines information is used with the global
  features of the material/molecule for property prediction tasks.
  
  Example
  -------
  >>> import deepchem as dc
  >>> from deepchem.models import MEGNetModel
  >>> from deepchem.utils.fake_data_generator import FakeGraphGenerator as FGG
  >>> graphs = FGG(global_features=4, num_classes=10).sample(n_graphs=20)
  >>> model = dc.models.MEGNetModel(n_node_features=5, n_edge_features=3, n_global_features=4, n_blocks=3, is_undirected=True, residual_connection=True, mode='classification', n_classes=10, batch_size=16)
  >>> training_loss = model.fit(graphs)
  
  References
  ----------
  .. [1] Chen, Chi, et al. "Graph networks as a universal machine learning framework for molecules and crystals." Chemistry of Materials 31.9 (2019): 3564-3572.
  .. [2] Battaglia, Peter W., et al. "Relational inductive biases, deep learning, and graph networks." arXiv preprint arXiv:1806.01261 (2018).
  Note
  ----
  The model requires PyTorch-Geometric to be installed.
docstr: Creates batch data for MEGNet model
  
  Note
  ----
  Ideally, we should only override default_generator method. But the problem
  here is that we _prepare_batch of TorchModel only supports non-graph
  data types. Hence, we are overriding it here. This should be fixed
  some time in the future.
docstr_tok: ["Creates","batch","data","for","MEGNet","model","Note","--","--","Ideally",",","we","should","only","override","default_generator","method",".","But","the","problem","here","is","that","we","_prepare_batch","of","TorchModel","only","supports","non-graph","data","types",".","Hence",",","we","are","overriding","it","here",".","This","should","be","fixed","some","time","in","the","future","."]
code: def _prepare_batch(self, batch):
      """Creates batch data for MEGNet model
  
      Note
      ----
      Ideally, we should only override default_generator method. But the problem
      here is that we _prepare_batch of TorchModel only supports non-graph
      data types. Hence, we are overriding it here. This should be fixed
      some time in the future.
      """
      try:
          from torch_geometric.data import Batch
      except ModuleNotFoundError:
          raise ImportError("This module requires PyTorch Geometric")
  
      # We convert deepchem.feat.GraphData to a PyG graph and then
      # batch it.
      graphs, labels, weights = batch
      # The default_generator method returns an array of dc.feat.GraphData objects
      # nested inside a list. To access the nested array of graphs, we are
      # indexing by 0 here.
      graph_list = [graph.to_pyg_graph() for graph in graphs[0]]
      pyg_batch = Batch()
      pyg_batch = pyg_batch.from_data_list(graph_list)
  
      _, labels, weights = super(MEGNetModel, self)._prepare_batch(
          ([], labels, weights)
      )
  
      return pyg_batch, labels, weights
  
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 MEGNetModel obj's underlying member method code with example usage: obj._prepare_batch(batch) and returns: pyg_batch,  labels,  weights