code: def assm(self, mol_batch, mol_tree_batch, mol_vec):
      cands = [
          mol_batch["cand_graph_batch"],
          cuda(mol_batch["tree_mess_src_e"]),
          cuda(mol_batch["tree_mess_tgt_e"]),
          cuda(mol_batch["tree_mess_tgt_n"]),
      ]
      cand_vec = self.jtmpn(cands, mol_tree_batch)
      cand_vec = self.G_mean(cand_vec)
  
      batch_idx = cuda(torch.LongTensor(mol_batch["cand_batch_idx"]))
      mol_vec = mol_vec[batch_idx]
  
      mol_vec = mol_vec.view(-1, 1, self.latent_size // 2)
      cand_vec = cand_vec.view(-1, self.latent_size // 2, 1)
      scores = (mol_vec @ cand_vec)[:, 0, 0]
  
      cnt, tot, acc = 0, 0, 0
      all_loss = []
      for i, mol_tree in enumerate(mol_batch["mol_trees"]):
          comp_nodes = [
              node_id
              for node_id, node in mol_tree.nodes_dict.items()
              if len(node["cands"]) > 1 and not node["is_leaf"]
          ]
          cnt += len(comp_nodes)
          # segmented accuracy and cross entropy
          for node_id in comp_nodes:
              node = mol_tree.nodes_dict[node_id]
              label = node["cands"].index(node["label"])
              ncand = len(node["cands"])
              cur_score = scores[tot : tot + ncand]
              tot += ncand
  
              if cur_score[label].item() >= cur_score.max().item():
                  acc += 1
  
              label = cuda(torch.LongTensor([label]))
              all_loss.append(
                  F.cross_entropy(
                      cur_score.view(1, -1), label, size_average=False
                  )
              )
  
      all_loss = sum(all_loss) / len(mol_batch["mol_trees"])
      return all_loss, acc / cnt
  
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 6 for calling the DGLJTNNVAE obj's underlying member method code with example usage: obj.assm(mol_batch, mol_tree_batch, mol_vec) and returns: all_loss,  unknown