code: def forward(self, mol_batch, beta=0, e1=None, e2=None):
      self.move_to_cuda(mol_batch)
  
      mol_trees = mol_batch["mol_trees"]
      batch_size = len(mol_trees)
  
      mol_tree_batch, tree_vec, mol_vec = self.encode(mol_batch)
  
      tree_vec, mol_vec, z_mean, z_log_var = self.sample(
          tree_vec, mol_vec, e1, e2
      )
      kl_loss = (
          -0.5
          * torch.sum(
              1.0 + z_log_var - z_mean * z_mean - torch.exp(z_log_var)
          )
          / batch_size
      )
  
      word_loss, topo_loss, word_acc, topo_acc = self.decoder(
          [t.graph for t in mol_trees], tree_vec
      )
      assm_loss, assm_acc = self.assm(
          mol_batch, mol_tree_batch, mol_vec
      )
      stereo_loss, stereo_acc = self.stereo(mol_batch, mol_vec)
  
      loss = (
          word_loss
          + topo_loss
          + assm_loss
          + 2 * stereo_loss
          + beta * kl_loss
      )
  
      return loss, kl_loss, word_acc, topo_acc, assm_acc, stereo_acc
  
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 5 for calling the DGLJTNNVAE obj's underlying member method code with example usage: obj.forward(mol_batch, beta, e1, e2) and returns: loss,  kl_loss,  word_acc,  topo_acc,  assm_acc,  stereo_acc