code: def encode(self, mol_batch):
      mol_graphs = mol_batch["mol_graph_batch"]
      mol_vec = self.mpn(mol_graphs)
  
      mol_tree_batch, tree_vec = self.jtnn(
          [t.graph for t in mol_batch["mol_trees"]]
      )
  
      self.n_nodes_total += mol_graphs.num_nodes()
      self.n_edges_total += mol_graphs.num_edges()
      self.n_tree_nodes_total += sum(
          t.graph.num_nodes() for t in mol_batch["mol_trees"]
      )
      self.n_passes += 1
  
      return mol_tree_batch, tree_vec, mol_vec
  
imports: import copy
  import rdkit.Chem
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import batch, unbatch
  from .chemutils import attach_mols_nx, copy_edit_mol, decode_stereo, enum_assemble_nx, set_atommap
  from .jtmpn import DGLJTMPN, mol2dgl_single
  from .jtnn_dec import DGLJTNNDecoder
  from .jtnn_enc import DGLJTNNEncoder
  from .mpn import DGLMPN, mol2dgl_single
  from .nnutils import cuda
label_desc: Use image node_id 3 for calling the DGLJTNNVAE obj's underlying member method code with example usage: obj.encode(mol_batch) and returns: mol_tree_batch,  tree_vec,  mol_vec