code: def run(self, mol_tree_batch, mol_tree_batch_lg, n_trees, tree_vec):
      node_offset = np.cumsum(
          np.insert(
              mol_tree_batch.batch_num_nodes().cpu().numpy(), 0, 0
          )
      )
      root_ids = node_offset[:-1]
      n_nodes = mol_tree_batch.num_nodes()
      n_edges = mol_tree_batch.num_edges()
  
      mol_tree_batch.ndata.update(
          {
              "x": self.embedding(mol_tree_batch.ndata["wid"]),
              "h": cuda(torch.zeros(n_nodes, self.hidden_size)),
              "new": cuda(
                  torch.ones(n_nodes).bool()
              ),  # whether it's newly generated node
          }
      )
  
      mol_tree_batch.edata.update(
          {
              "s": cuda(torch.zeros(n_edges, self.hidden_size)),
              "m": cuda(torch.zeros(n_edges, self.hidden_size)),
              "r": cuda(torch.zeros(n_edges, self.hidden_size)),
              "z": cuda(torch.zeros(n_edges, self.hidden_size)),
              "src_x": cuda(torch.zeros(n_edges, self.hidden_size)),
              "dst_x": cuda(torch.zeros(n_edges, self.hidden_size)),
              "rm": cuda(torch.zeros(n_edges, self.hidden_size)),
              "accum_rm": cuda(torch.zeros(n_edges, self.hidden_size)),
          }
      )
  
      mol_tree_batch.apply_edges(
          func=lambda edges: {
              "src_x": edges.src["x"],
              "dst_x": edges.dst["x"],
          },
      )
  
      # input tensors for stop prediction (p) and label prediction (q)
      p_inputs = []
      p_targets = []
      q_inputs = []
      q_targets = []
  
      # Predict root
      mol_tree_batch.pull(
          root_ids, DGLF.copy_e("m", "m"), DGLF.sum("m", "h")
      )
      mol_tree_batch.apply_nodes(dec_tree_node_update, v=root_ids)
      # Extract hidden states and store them for stop/label prediction
      h = mol_tree_batch.nodes[root_ids].data["h"]
      x = mol_tree_batch.nodes[root_ids].data["x"]
      p_inputs.append(torch.cat([x, h, tree_vec], 1))
      # If the out degree is 0 we don't generate any edges at all
      root_out_degrees = mol_tree_batch.out_degrees(root_ids)
      q_inputs.append(torch.cat([h, tree_vec], 1))
      q_targets.append(mol_tree_batch.nodes[root_ids].data["wid"])
  
      # Traverse the tree and predict on children
      for eid, p in dfs_order(mol_tree_batch, root_ids):
          eid = eid.to(mol_tree_batch.device)
          p = p.to(mol_tree_batch.device)
          u, v = mol_tree_batch.find_edges(eid)
  
          p_target_list = torch.zeros_like(root_out_degrees)
          p_target_list[root_out_degrees > 0] = 1 - p
          p_target_list = p_target_list[root_out_degrees >= 0]
          p_targets.append(torch.tensor(p_target_list))
  
          root_out_degrees -= (root_out_degrees == 0).long()
          root_out_degrees -= torch.tensor(
              np.isin(root_ids, v.cpu().numpy())
          ).to(root_out_degrees)
  
          mol_tree_batch_lg.ndata.update(mol_tree_batch.edata)
          mol_tree_batch_lg.pull(
              eid, DGLF.copy_u("m", "m"), DGLF.sum("m", "s")
          )
          mol_tree_batch_lg.pull(
              eid, DGLF.copy_u("rm", "rm"), DGLF.sum("rm", "accum_rm")
          )
          mol_tree_batch_lg.apply_nodes(
              self.dec_tree_edge_update, v=eid
          )
          mol_tree_batch.edata.update(mol_tree_batch_lg.ndata)
  
          is_new = mol_tree_batch.nodes[v].data["new"]
          mol_tree_batch.pull(
              v, DGLF.copy_e("m", "m"), DGLF.sum("m", "h")
          )
          mol_tree_batch.apply_nodes(dec_tree_node_update, v=v)
  
          # Extract
          n_repr = mol_tree_batch.nodes[v].data
          h = n_repr["h"]
          x = n_repr["x"]
          tree_vec_set = tree_vec[root_out_degrees >= 0]
          wid = n_repr["wid"]
          p_inputs.append(torch.cat([x, h, tree_vec_set], 1))
          # Only newly generated nodes are needed for label prediction
          # NOTE: The following works since the uncomputed messages are zeros.
  
          q_input = torch.cat([h, tree_vec_set], 1)[is_new]
          q_target = wid[is_new]
          if q_input.shape[0] > 0:
              q_inputs.append(q_input)
              q_targets.append(q_target)
      p_targets.append(
          torch.zeros(
              (root_out_degrees == 0).sum(),
              device=root_out_degrees.device,
              dtype=torch.int64,
          )
      )
  
      # Batch compute the stop/label prediction losses
      p_inputs = torch.cat(p_inputs, 0)
      p_targets = cuda(torch.cat(p_targets, 0))
      q_inputs = torch.cat(q_inputs, 0)
      q_targets = torch.cat(q_targets, 0)
  
      q = self.W_o(torch.relu(self.W(q_inputs)))
      p = self.U_s(torch.relu(self.U(p_inputs)))[:, 0]
  
      p_loss = (
          F.binary_cross_entropy_with_logits(
              p, p_targets.float(), size_average=False
          )
          / n_trees
      )
      q_loss = (
          F.cross_entropy(q, q_targets, size_average=False) / n_trees
      )
      p_acc = (
          (p > 0).long() == p_targets
      ).sum().float() / p_targets.shape[0]
      q_acc = (
          q.max(1)[1] == q_targets
      ).float().sum() / q_targets.shape[0]
  
      self.q_inputs = q_inputs
      self.q_targets = q_targets
      self.q = q
      self.p_inputs = p_inputs
      self.p_targets = p_targets
      self.p = p
  
      return q_loss, p_loss, q_acc, p_acc
  
file_path: dgl/examples/pytorch/jtnn/jtnn/jtnn_dec.py
imports: import dgl.function
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import batch, dfs_labeled_edges_generator, line_graph
  from .chemutils import enum_assemble_nx, get_mol
  from .mol_tree_nx import DGLMolTree
  from .nnutils import cuda, GRUUpdate, tocpu
label_desc: Use image node_id 3 for calling the DGLJTNNDecoder obj's underlying member method code with example usage: obj.run(mol_tree_batch, mol_tree_batch_lg, n_trees, tree_vec) and returns: q_loss,  p_loss,  q_acc,  p_acc