code: def test():
      dataset.training = False
      dataloader = DataLoader(
          dataset,
          batch_size=1,
          shuffle=False,
          num_workers=0,
          collate_fn=JTNNCollator(vocab, False),
          drop_last=True,
          worker_init_fn=worker_init_fn,
      )
  
      # Just an example of molecule decoding; in reality you may want to sample
      # tree and molecule vectors.
      for it, batch in enumerate(dataloader):
          gt_smiles = batch["mol_trees"][0].smiles
          print(gt_smiles)
          model.move_to_cuda(batch)
          _, tree_vec, mol_vec = model.encode(batch)
          tree_vec, mol_vec, _, _ = model.sample(tree_vec, mol_vec)
          smiles = model.decode(tree_vec, mol_vec)
          print(smiles)
  
file_path: dgl/examples/pytorch/jtnn/vaetrain_dgl.py
imports: import math
  import random
  import sys
  from collections import deque
  from optparse import OptionParser
  import rdkit
  import torch
  import torch.nn
  import torch.optim
  import torch.optim.lr_scheduler
  import tqdm
  from jtnn import *
  from torch.utils.data import DataLoader
label_desc: Use image node_id 3 for calling a global function with example usage: test() without return types