Repository: awslabs/dgl-lifesci (language: python)

Issue:
JTNN example does not work
Hi,

I am trying to run the JTNN example in dgl-lifesci/examples/generative_models/jtvae. When I run python train.py, I get this error

"(rdkit-env) root@nucar-nice:/home/trinayan/dgl-lifesci/examples/generative_models/jtvae# python3 train.py 
Using backend: pytorch
Traceback (most recent call last):
  File "train.py", line 133, in <module>
    main(args)
  File "train.py", line 31, in main
    depth=args.depth)
  File "/root/anaconda3/envs/rdkit-env/lib/python3.6/site-packages/dgllife-0.2.5-py3.6.egg/dgllife/model/model_zoo/jtnn/jtnn_vae.py", line 53, in __init__
FileNotFoundError: [Errno 2] No such file or directory: '<built-in function dir>/jtnn/vocab.txt'"

Not sure if i need to download anything since the github page says the datasets will be downloaded automatically for ZINC. 

Any help will be appreciated

Thanks


Discussion / hints:
Has anyone had a chance to look into this? . Thanks
Hi @trinayan, sorry for the inconvenience. I think the training part of the JTNN example is broken due to DGL 0.5 and I'm working on fixing it, which might take some time. If you are in an urgent need of using that, I would recommend you to directly use the [authors' code](https://github.com/wengong-jin/icml18-jtnn).

Fix patch:
diff --git a/docs/source/api/data.rst b/docs/source/api/data.rst
index 6feaaead..40435c7a 100644
--- a/docs/source/api/data.rst
+++ b/docs/source/api/data.rst
@@ -148,21 +148,6 @@ Adapting to New Datasets for Weisfeiler-Lehman Networks
 .. autoclass:: dgllife.data.WLNRankDataset
     :members: ignore_large, __getitem__, __len__
 
-Generative Models
------------------
-
-JTVAE
-`````
-
-.. autoclass:: dgllife.data.DGLMolTree
-    :members: treesize, recover, assemble
-
-.. autoclass:: dgllife.data.JTVAEDataset
-    :members: __len__, __getitem__, move_to_device
-
-.. autoclass:: dgllife.data.JTVAECollator
-    :members: __call__
-
 Protein-Ligand Binding Affinity Prediction
 ------------------------------------------
 
diff --git a/docs/source/api/model.zoo.rst b/docs/source/api/model.zoo.rst
index bd7cf1e7..b4583dab 100644
--- a/docs/source/api/model.zoo.rst
+++ b/docs/source/api/model.zoo.rst
@@ -77,11 +77,6 @@ DGMG
 .. automodule:: dgllife.model.model_zoo.dgmg
     :members:
 
-JTNN
-````
-.. autoclass:: dgllife.model.model_zoo.jtnn.DGLJTNNVAE
-    :members:
-
 Reaction Prediction
 
 WLN for Reaction Center Prediction
diff --git a/examples/generative_models/jtvae/README.md b/examples/generative_models/jtvae/README.md
index 860d8169..81045085 100644
--- a/examples/generative_models/jtvae/README.md
+++ b/examples/generative_models/jtvae/README.md
@@ -1,129 +1,28 @@
-# Junction Tree Variational Autoencoder for Molecular Graph Generation (JTNN)
+# Junction Tree Variational Autoencoder for Molecular Graph Generation
 
-Wengong Jin, Regina Barzilay, Tommi Jaakkola. 
-Junction Tree Variational Autoencoder for Molecular Graph Generation. 
-*arXiv preprint arXiv:1802.04364*, 2018.
+- [paper](https://arxiv.org/abs/1802.04364)
+- [authors' code](https://github.com/wengong-jin/icml18-jtnn/tree/master/molvae)
 
-JTNN uses algorithm called junction tree algorithm to form a tree from the molecular graph. 
-Then the model will encode the tree and graph into two separate vectors `z_G` and `z_T`. Details can
-be found in original paper. The brief process is as below (from original paper): 
+## Training
 
-![image](https://user-images.githubusercontent.com/8686776/63677300-3fb6d980-c81f-11e9-8a65-57c8b03aaf52.png)
+We trained VAE model in two phases:
 
-**Goal**: JTNN is an auto-encoder model, aiming to learn hidden representation for molecular graphs. 
-These representations can be used for downstream tasks, such as property prediction, or molecule optimizations.
-
-## Dataset
-
-### ZINC
-
-> The ZINC database is a curated collection of commercially available chemical compounds 
-prepared especially for virtual screening. (introduction from Wikipedia)
-
-Generally speaking, molecules in the ZINC dataset are more drug-like. We use ~220,000 
-molecules for training and 5000 molecules for validation. 
-
-### GuacaMol
-
-> The dataset is used by the [GuacaMol](https://github.com/BenevolentAI/guacamol) benchmark of generative models.  The dataset is derived from the [ChEMBL](https://www.ebi.ac.uk/chembl/) 24 database.
-
-Generally speaking, the advantage of ChEMBL is that it contains only molecules that have been synthesized and
-tested against biological targets. Other datasets, such as ZINC, contain, at least to some
-degree, virtual molecules that are likely to be synthesizable, but have not been made yet. (introduction from their [paper](https://arxiv.org/abs/1811.09621)) 
-We use ~1,270,000 molecules for training and ~50,000 molecules for validation.
-
-### Download
-For ZINC dataset, you can run with
-```
-python train.py
-```
-The dataset will be downloaded automatically. 
-
-For GuacaMol dataset, you can download the training-set and validation-set with 
-```
-mkdir -p $DIRECTORY
-wget https://ndownloader.figshare.com/files/13612760 -O $DIRECTORY/guacamol_v1_train.smiles
-wget https://ndownloader.figshare.com/files/13612766 -O $DIRECTORY/guacamol_v1_valid.smiles
-``` 
-`$DIRECTORY` specifies the directory to save the dataset files.
-
-## Usage
-
-### Training
-#### ZINC
-For ZINC dataset, to start training, use `python train.py`. By default, the script will use ZINC dataset
- with preprocessed vocabulary, and save model checkpoint periodically in the current working directory. 
- 
-#### GuacaMol
-For GuacaMol dataset, you need to specify the downloaded training-set file and the preprocessed vocabulary. To start training, use
-```
-python train.py -t $DIRECTORY/guacamol_v1_train.smiles -v guacamol -s $CHECKPOINTS
-```
-`$DIRECTORY` specifies the directory of downloaded dataset files.  
-`$CHECKPOINTS` specifies the directory to save model checkpoint periodically. 
-
-
-### Evaluation
-
-To start evaluation, use `python reconstruct_eval.py`. By default, we will perform evaluation with 
-DGL's pre-trained model. During the evaluation, the program will print out the success rate of 
-molecule reconstruction.
-
-### Pre-trained models
-
-Below gives the statistics of our pre-trained `JTNN_ZINC` model. 
-
-| Pre-trained model  | % Reconstruction Accuracy
-| ------------------ | -------
-| `JTNN_ZINC`        |  73.7             
-
-### Visualization
-
-Here we draw some "neighbor" of a given molecule, by adding noises on the intermediate representations. 
-You can download the script with `wget https://data.dgl.ai/dgllife/jtnn_viz_neighbor_mol.ipynb`. 
-Please put this script at the current directory (`examples/pytorch/model_zoo/chem/generative_models/jtnn/`).
-
-#### Given Molecule
-![image](https://user-images.githubusercontent.com/8686776/63773593-0d37da00-c90e-11e9-8933-0abca4b430db.png)
-#### Neighbor Molecules
-![image](https://user-images.githubusercontent.com/8686776/63773602-1163f780-c90e-11e9-8341-5122dc0d0c82.png)
-
-### Dataset configuration
-
-If you want to use your own dataset, please create a file with one SMILES a line as below
+1. We first train a model without KL regularization term. The model checkpoints will be saved as `pre_model/model.iter-x`, where `x + 1` is the number of training epochs.
+2. We then train the model with KL regularization by passing the path to a saved model checkpoint. The model checkpoints will be saved as `vae_model/model.iter-x`.
 
-```
-CCO
-Fc1ccccc1
+```bash
+CUDA_VISIBLE_DEVICES=0 python pretrain.py
+CUDA_VISIBLE_DEVICES=0 python vaetrain.py -m pre_model/model.iter-2
 ```
 
-You can generate the vocabulary file corresponding to your dataset with `python vocab.py -d X -v Y`, where `X` 
-is the path to the dataset and `Y` is the path to the vocabulary file to save. An example vocabulary file 
-corresponding to the two molecules above will be
+Note that the weight of the KL regularization term generally controls a trade off between reconstruction accuracy and generation diversity. To adjust the weight, specify `-z` (default: 0.001).
 
-```
-CC
-CF
-C1=CC=CC=C1
-CO
-```
+## Testing
 
-If you want to develop a model based on DGL's pre-trained model, it's important to make sure that the vocabulary 
-generated above is a subset of the vocabulary we use for the pre-trained model. By running `vocab.py` above, we 
-also check if the new vocabulary is a subset of the vocabulary we use for the pre-trained model and print the 
-result in the terminal as follows:
+For molecule reconstruction, 
 
+```bash
+CUDA_VISIBLE_DEVICES=0 python reconstruct.py -m Y
 ```
-The new vocabulary is a subset of the default vocabulary: True
-```
-
-To train on this new dataset, run
-
-```
-python train.py -t X
-```
-
-where `X` is the path to the new dataset. If you want to use the vocabulary generated above, also add `-v Y`, where 
-`Y` is the path to the vocabulary file we just saved.
 
-To evaluate on this new dataset, run `python reconstruct_eval.py` with arguments same as above.
+where `Y` is the path to a model checkpoint. If not specified, this will evaluate on a pre-trained model trained without KL regularization.
diff --git a/examples/generative_models/jtvae/pretrain.py b/examples/generative_models/jtvae/pretrain.py
new file mode 100644
index 00000000..2641c35c
--- /dev/null
+++ b/examples/generative_models/jtvae/pretrain.py
@@ -0,0 +1,131 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# Pre-trained an AE
+
+import rdkit
+import numpy as np
+import sys
+import time
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+from dgllife.data import JTVAEZINC, JTVAEDataset, JTVAECollator
+from dgllife.model import JTNNVAE
+from dgllife.utils import JTVAEVocab
+from torch.utils.data import DataLoader
+
+from utils import mkdir_p
+
+def main(args):
+    mkdir_p(args.save_path)
+
+    lg = rdkit.RDLogger.logger()
+    lg.setLevel(rdkit.RDLogger.CRITICAL)
+
+    if args.use_cpu or not torch.cuda.is_available():
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda:0')
+
+    vocab = JTVAEVocab(file_path=args.train_path)
+    if args.train_path is None:
+        dataset = JTVAEZINC('train', vocab)
+    else:
+        dataset = JTVAEDataset(args.train_path, vocab, training=True)
+    dataloader = DataLoader(dataset,
+                            batch_size=args.batch_size,
+                            shuffle=True,
+                            num_workers=args.num_workers,
+                            collate_fn=JTVAECollator(training=True),
+                            drop_last=True)
+
+    model = JTNNVAE(vocab, args.hidden_size, args.latent_size, args.depth).to(device)
+    model.reset_parameters()
+    print("Model #Params: {:d}K".format(sum([x.nelement() for x in model.parameters()]) // 1000))
+
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+    scheduler = lr_scheduler.ExponentialLR(optimizer, args.gamma)
+
+    dur = []
+    t0 = time.time()
+    for epoch in range(args.max_epoch):
+        word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
+
+        for it, (batch_trees, batch_tree_graphs, batch_mol_graphs,
+                 stereo_cand_batch_idx, stereo_cand_labels, batch_stereo_cand_graphs) \
+                in enumerate(dataloader):
+            batch_tree_graphs = batch_tree_graphs.to(device)
+            batch_mol_graphs = batch_mol_graphs.to(device)
+            stereo_cand_batch_idx = stereo_cand_batch_idx.to(device)
+            batch_stereo_cand_graphs = batch_stereo_cand_graphs.to(device)
+
+            loss, kl_div, wacc, tacc, sacc, dacc = model(
+                batch_trees, batch_tree_graphs, batch_mol_graphs, stereo_cand_batch_idx,
+                stereo_cand_labels, batch_stereo_cand_graphs, beta=0)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            word_acc += wacc
+            topo_acc += tacc
+            assm_acc += sacc
+            steo_acc += dacc
+
+            if (it + 1) % args.print_iter == 0:
+                dur.append(time.time() - t0)
+                word_acc = word_acc / args.print_iter * 100
+                topo_acc = topo_acc / args.print_iter * 100
+                assm_acc = assm_acc / args.print_iter * 100
+                steo_acc = steo_acc / args.print_iter * 100
+
+                print('Epoch {:d}/{:d} | Iter {:d}/{:d} | KL: {:.1f}, Word: {:.2f}, '
+                      'Topo: {:.2f}, Assm: {:.2f}, Steo: {:.2f} | '
+                      'Estimated time per epoch: {:.4f}'.format(
+                    epoch + 1, args.max_epoch, it + 1, len(dataloader), kl_div, word_acc,
+                    topo_acc, assm_acc, steo_acc, np.mean(dur) / args.print_iter * len(dataloader)))
+                word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
+                sys.stdout.flush()
+                t0 = time.time()
+
+        scheduler.step()
+        print("learning rate: {:.6f}".format(scheduler.get_lr()[0]))
+        torch.save(model.state_dict(), args.save_path + "/model.iter-" + str(epoch))
+
+if __name__ == '__main__':
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument('-tr', '--train-path', type=str,
+                        help='Path to the training molecules, with one SMILES string a line')
+    parser.add_argument('-s', '--save-path', type=str, default='pre_model',
+                        help='Directory to save model checkpoints')
+    parser.add_argument('-b', '--batch-size', type=int, default=40,
+                        help='Batch size')
+    parser.add_argument('-w', '--hidden-size', type=int, default=450,
+                        help='Hidden size')
+    parser.add_argument('-l', '--latent-size', type=int, default=56,
+                        help='Latent size')
+    parser.add_argument('-d', '--depth', type=int, default=3,
+                        help='Number of GNN layers')
+    parser.add_argument('-lr', '--lr', type=float, default=1e-3,
+                        help='Learning rate')
+    parser.add_argument('-g', '--gamma', type=float, default=0.9,
+                        help='Multiplicative factor for learning rate decay')
+    parser.add_argument('-me', '--max-epoch', type=int, default=3,
+                        help='Maximum number of epochs for training')
+    parser.add_argument('-nw', '--num-workers', type=int, default=4,
+                        help='Number of subprocesses for data loading')
+    parser.add_argument('-pi', '--print-iter', type=int, default=20,
+                        help='Frequency for printing evaluation metrics')
+    parser.add_argument('-cpu', '--use-cpu', action='store_true',
+                        help='By default, the script uses GPU whenever available. '
+                             'This flag enforces the use of CPU.')
+    args = parser.parse_args()
+
+    main(args)
diff --git a/examples/generative_models/jtvae/reconstruct.py b/examples/generative_models/jtvae/reconstruct.py
new file mode 100644
index 00000000..6df41eb1
--- /dev/null
+++ b/examples/generative_models/jtvae/reconstruct.py
@@ -0,0 +1,71 @@
+import rdkit
+import torch
+
+from dgllife.data import JTVAEZINC, JTVAEDataset, JTVAECollator
+from dgllife.utils import JTVAEVocab
+from dgllife.model import JTNNVAE, load_pretrained
+from torch.utils.data import DataLoader
+
+def main(args):
+    lg = rdkit.RDLogger.logger()
+    lg.setLevel(rdkit.RDLogger.CRITICAL)
+
+    if args.use_cpu or not torch.cuda.is_available():
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda:0')
+
+    vocab = JTVAEVocab(file_path=args.train_path)
+    if args.test_path is None:
+        dataset = JTVAEZINC('test', vocab)
+    else:
+        dataset = JTVAEDataset(args.test_path, vocab, training=False)
+    dataloader = DataLoader(dataset,
+                            batch_size=1,
+                            collate_fn=JTVAECollator(training=False))
+
+    if args.model_path is None:
+        model = load_pretrained('JTVAE_ZINC_no_kl')
+    else:
+        model = JTNNVAE(vocab, args.hidden_size, args.latent_size, args.depth)
+        model.load_state_dict(torch.load(args.model_path, map_location='cpu'))
+    model = model.to(device)
+
+    acc = 0.0
+    for it, (tree, tree_graph, mol_graph) in enumerate(dataloader):
+        tot = it + 1
+        smiles = tree.smiles
+        tree_graph = tree_graph.to(device)
+        mol_graph = mol_graph.to(device)
+        dec_smiles = model.reconstruct(tree_graph, mol_graph)
+        if dec_smiles == smiles:
+            acc += 1
+        if tot % args.print_iter == 0:
+            print('Iter {:d}/{:d} | Acc {:.4f}'.format(
+                tot // args.print_iter, len(dataloader) // args.print_iter, acc / tot))
+    print('Final acc: {:.4f}'.format(acc / tot))
+
+if __name__ == '__main__':
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument('-tr', '--train-path', type=str,
+                        help='Path to the training molecules, with one SMILES string a line')
+    parser.add_argument('-te', '--test-path', type=str,
+                        help='Path to the test molecules, with one SMILES string a line')
+    parser.add_argument('-m', '--model-path', type=str,
+                        help='Path to pre-trained model checkpoint')
+    parser.add_argument('-w', '--hidden-size', type=int, default=450,
+                        help='Hidden size')
+    parser.add_argument('-l', '--latent-size', type=int, default=56,
+                        help='Latent size')
+    parser.add_argument('-d', '--depth', type=int, default=3,
+                        help='Number of GNN layers')
+    parser.add_argument('-pi', '--print-iter', type=int, default=20,
+                        help='Frequency for printing evaluation metrics')
+    parser.add_argument('-cpu', '--use-cpu', action='store_true',
+                        help='By default, the script uses GPU whenever available. '
+                             'This flag enforces the use of CPU.')
+    args = parser.parse_args()
+
+    main(args)
diff --git a/examples/generative_models/jtvae/reconstruct_eval.py b/examples/generative_models/jtvae/reconstruct_eval.py
deleted file mode 100644
index b14882bc..00000000
--- a/examples/generative_models/jtvae/reconstruct_eval.py
+++ /dev/null
@@ -1,108 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import torch
-
-from dgllife.data import JTVAEDataset, JTVAECollator
-from dgllife.model import DGLJTNNVAE, load_pretrained
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-from utils import worker_init_fn, get_vocab_file
-
-@torch.no_grad()
-def main(args):
-    worker_init_fn(None)
-
-    if torch.cuda.is_available():
-        device = torch.device('cuda:0')
-    else:
-        device = torch.device('cpu')
-
-    if args.model_path is not None:
-        model = DGLJTNNVAE(vocab_file=get_vocab_file(args.vocab),
-                           hidden_size=args.hidden_size,
-                           latent_size=args.latent_size,
-                           depth=args.depth)
-        model.load_state_dict(torch.load(args.model_path))
-    else:
-        model = load_pretrained("JTNN_ZINC")
-    print("# model parameters: {:d}K".format(
-        sum([x.nelement() for x in model.parameters()]) // 1000))
-
-    dataset = JTVAEDataset(data=args.data, vocab=model.vocab, training=False)
-    dataloader = DataLoader(
-        dataset,
-        batch_size=1,
-        shuffle=False,
-        num_workers=0,
-        collate_fn=JTVAECollator(False),
-        worker_init_fn=worker_init_fn)
-    # Just an example of molecule decoding; in reality you may want to sample
-    # tree and molecule vectors.
-    acc = 0.0
-    tot = 0
-    model = model.to(device)
-    model.eval()
-    for it, batch in enumerate(tqdm(dataloader)):
-        gt_smiles = batch['mol_trees'][0].smiles
-        batch = dataset.move_to_device(batch, device)
-        try:
-            _, tree_vec, mol_vec = model.encode(batch)
-
-            tree_mean = model.T_mean(tree_vec)
-            # Following Mueller et al.
-            tree_log_var = -torch.abs(model.T_var(tree_vec))
-            epsilon = torch.randn(1, model.latent_size // 2).to(device)
-            tree_vec = tree_mean + torch.exp(tree_log_var // 2) * epsilon
-
-            mol_mean = model.G_mean(mol_vec)
-            # Following Mueller et al.
-            mol_log_var = -torch.abs(model.G_var(mol_vec))
-            epsilon = torch.randn(1, model.latent_size // 2).to(device)
-            mol_vec = mol_mean + torch.exp(mol_log_var // 2) * epsilon
-
-            dec_smiles = model.decode(tree_vec, mol_vec)
-
-            if dec_smiles == gt_smiles:
-                acc += 1
-            tot += 1
-        except Exception as e:
-            print("Failed to encode: {}".format(gt_smiles))
-            print(e)
-
-        if it % 20 == 1:
-            print("Progress {}/{}; Current Reconstruction Accuracy: {:.4f}".format(
-                it, len(dataloader), acc / tot))
-
-    print("Reconstruction Accuracy: {}".format(acc / tot))
-
-if __name__ == '__main__':
-    from argparse import ArgumentParser
-
-    parser = ArgumentParser(description="Evaluation for JTVAE")
-
-    parser.add_argument("-da", "--data", default='test',
-                        help="Data for evaluation. 1) If 'train' or 'test', it will perform the "
-                             "evaluation on the training or test subset of ZINC. 2) Otherwise, it "
-                             "should be the path to a .txt file of compounds, with one SMILES per line.")
-    parser.add_argument("-v", "--vocab", default='zinc',
-                        help="Vocabulary to use. 1) If 'zinc' or 'guacamol', it will use the "
-                             "vocabulary extracted from ZINC or GuacaMol. 2) Otherwise, it "
-                             "should be the path to a file of vocabulary generated with vocab.py")
-    parser.add_argument("-m", "--model-path", default=None,
-                        help="Path to a pre-trained model. If not specified, it will load a model "
-                             "pre-trained on ZINC.")
-    parser.add_argument("-w", "--hidden-size", type=int, default=450,
-                        help="Hidden size of representation vector, "
-                             "should be consistent with pre-trained model")
-    parser.add_argument("-l", "--latent-size", type=int, default=56,
-                        help="Latent size of node features and edge features, "
-                             "should be consistent with pre-trained model")
-    parser.add_argument("-d", "--depth", type=int, default=3,
-                        help="Depth of message passing hops, "
-                             "should be consistent with pre-trained model")
-    args = parser.parse_args()
-    main(args)
diff --git a/examples/generative_models/jtvae/train.py b/examples/generative_models/jtvae/train.py
deleted file mode 100644
index 6e63dc84..00000000
--- a/examples/generative_models/jtvae/train.py
+++ /dev/null
@@ -1,133 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import sys
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torch.optim.lr_scheduler as lr_scheduler
-
-from dgllife.data import JTVAEDataset, JTVAECollator
-from dgllife.model import DGLJTNNVAE
-from torch.utils.data import DataLoader
-
-from utils import worker_init_fn, get_vocab_file
-
-def main(args):
-    torch.multiprocessing.set_sharing_strategy('file_system')
-
-    worker_init_fn(None)
-
-    if torch.cuda.is_available():
-        device = torch.device('cuda:0')
-    else:
-        device = torch.device('cpu')
-
-    model = DGLJTNNVAE(vocab_file=get_vocab_file(args.vocab),
-                       hidden_size=args.hidden_size,
-                       latent_size=args.latent_size,
-                       depth=args.depth)
-    print("# model parameters: {:d}K".format(
-        sum([x.nelement() for x in model.parameters()]) // 1000))
-
-    if args.model_path is not None:
-        model.load_state_dict(torch.load(args.model_path))
-    else:
-        for param in model.parameters():
-            if param.dim() == 1:
-                nn.init.constant_(param, 0)
-            else:
-                nn.init.xavier_normal_(param)
-
-    dataset = JTVAEDataset(data=args.data, vocab=model.vocab, training=True)
-    dataloader = DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        num_workers=4,
-        collate_fn=JTVAECollator(True),
-        drop_last=True,
-        worker_init_fn=worker_init_fn)
-
-    optimizer = optim.Adam(model.parameters(), lr=args.lr)
-    scheduler = lr_scheduler.ExponentialLR(optimizer, 0.9)
-
-    model = model.to(device)
-    model.train()
-    for epoch in range(args.max_epoch):
-        word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
-
-        for it, batch in enumerate(dataloader):
-            batch = dataset.move_to_device(batch, device)
-            model.zero_grad()
-            try:
-                loss, kl_div, wacc, tacc, sacc, dacc = model(batch, args.beta)
-            except:
-                print([t.smiles for t in batch['mol_trees']])
-                raise
-            loss.backward()
-            optimizer.step()
-
-            word_acc += wacc
-            topo_acc += tacc
-            assm_acc += sacc
-            steo_acc += dacc
-
-            if (it + 1) % args.print_iter == 0:
-                word_acc = word_acc / args.print_iter * 100
-                topo_acc = topo_acc / args.print_iter * 100
-                assm_acc = assm_acc / args.print_iter * 100
-                steo_acc = steo_acc / args.print_iter * 100
-
-                print("KL: %.1f, Word: %.2f, Topo: %.2f, Assm: %.2f, Steo: %.2f, Loss: %.6f" % (
-                    kl_div, word_acc, topo_acc, assm_acc, steo_acc, loss.item()))
-                word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
-                sys.stdout.flush()
-
-            if (it + 1) % 1500 == 0:  # Fast annealing
-                scheduler.step()
-                print("learning rate: %.6f" % scheduler.get_lr()[0])
-                torch.save(model.state_dict(),
-                           args.save_path + "/model.iter-%d-%d" % (epoch, it + 1))
-
-        scheduler.step()
-        print("learning rate: %.6f" % scheduler.get_lr()[0])
-        torch.save(model.state_dict(), args.save_path + "/model.iter-" + str(epoch))
-
-if __name__ == '__main__':
-    from argparse import ArgumentParser
-
-    parser = ArgumentParser(description="Training for JTVAE")
-    parser.add_argument("-da", "--data", default='train',
-                        help="Data for training. 1) If 'train', it will perform the training on the "
-                             "training subset of ZINC. 2) Otherwise, it should be the path to a .txt "
-                             "file of compounds, with one SMILES per line.")
-    parser.add_argument("-v", "--vocab", default='zinc',
-                        help="Vocabulary to use. 1) If 'zinc' or 'guacamol', it will use the "
-                             "vocabulary extracted from ZINC or GuacaMol. 2) Otherwise, it "
-                             "should be the path to a file of vocabulary generated with vocab.py")
-    parser.add_argument("-s", "--save_dir", dest="save_path", default='./',
-                        help="Path to save checkpoint models, default to be current working directory")
-    parser.add_argument("-m", "--model", dest="model_path", default=None,
-                        help="Path to load pre-trained model")
-    parser.add_argument("-b", "--batch-size", type=int, default=40,
-                        help="Batch size")
-    parser.add_argument("-w", "--hidden-size", type=int, default=200,
-                        help="Size of representation vectors")
-    parser.add_argument("-l", "--latent-size", type=int, default=56,
-                        help="Latent Size of node(atom) features and edge(atom) features")
-    parser.add_argument("-d", "--depth", type=int, default=3,
-                        help="Depth of message passing hops")
-    parser.add_argument("-z", "--beta", type=float, default=1.0,
-                        help="Coefficient of KL Divergence term")
-    parser.add_argument("-q", "--lr", type=float, default=1e-3,
-                        help="Learning Rate")
-    parser.add_argument("-me", "--max-epoch", type=int, default=100,
-                        help='Max number of epochs for training')
-    parser.add_argument("-pi", "--print-iter", type=int, default=20,
-                        help='Frequency of printing messages')
-    args = parser.parse_args()
-
-    main(args)
diff --git a/examples/generative_models/jtvae/utils.py b/examples/generative_models/jtvae/utils.py
index 190cca2b..492eab27 100644
--- a/examples/generative_models/jtvae/utils.py
+++ b/examples/generative_models/jtvae/utils.py
@@ -2,33 +2,27 @@
 #
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# Utils for JTVAE
 
-import rdkit
-
-def worker_init_fn(id_):
-    lg = rdkit.RDLogger.logger()
-    lg.setLevel(rdkit.RDLogger.CRITICAL)
+import errno
+import os
 
-def get_vocab_file(vocab):
-    """Get the path to the vocabulary file.
+def mkdir_p(path):
+    """Create a folder for the given path.
 
     Parameters
     ----------
-    vocab : str
-        * If 'zinc' or 'guacamol', it will return the path to the
-          vocabulary file for the corresponding dataset.
-        * Otherwise, it should be already the path to a vocabulary file.
-
-    Returns
-    -------
-    str
-        The path to a vocabulary file.
+    path: str
+        Folder to create
     """
-    if vocab == 'zinc':
-        vocab_file = '{}/jtnn/vocab.txt'.format(dir)
-    elif vocab == 'guacamol':
-        vocab_file = '{}/jtnn/vocab_guacamol.txt'.format(dir)
-    else:
-        vocab_file = vocab
-
-    return vocab_file
+    try:
+        os.makedirs(path)
+        print('Created directory {}'.format(path))
+    except OSError as exc:
+        if exc.errno == errno.EEXIST and os.path.isdir(path):
+            print('Directory {} already exists.'.format(path))
+        else:
+            raise
diff --git a/examples/generative_models/jtvae/vaetrain.py b/examples/generative_models/jtvae/vaetrain.py
new file mode 100644
index 00000000..9a522dd5
--- /dev/null
+++ b/examples/generative_models/jtvae/vaetrain.py
@@ -0,0 +1,134 @@
+import numpy as np
+import rdkit
+import sys
+import time
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+from dgllife.data import JTVAEZINC, JTVAEDataset, JTVAECollator
+from dgllife.utils import JTVAEVocab
+from dgllife.model import JTNNVAE
+from torch.utils.data import DataLoader
+
+from utils import mkdir_p
+
+def main(args):
+    mkdir_p(args.save_path)
+
+    lg = rdkit.RDLogger.logger()
+    lg.setLevel(rdkit.RDLogger.CRITICAL)
+
+    if args.use_cpu or not torch.cuda.is_available():
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda:0')
+
+    vocab = JTVAEVocab(file_path=args.train_path)
+    if args.train_path is None:
+        dataset = JTVAEZINC('train', vocab)
+    else:
+        dataset = JTVAEDataset(args.train_path, vocab, training=True)
+    dataloader = DataLoader(dataset,
+                            batch_size=args.batch_size,
+                            shuffle=True,
+                            num_workers=args.num_workers,
+                            collate_fn=JTVAECollator(training=True),
+                            drop_last=True)
+
+    model = JTNNVAE(vocab, args.hidden_size, args.latent_size, args.depth)
+    if args.model_path is not None:
+        model.load_state_dict(torch.load(args.model_path, map_location='cpu'))
+    else:
+        model.reset_parameters()
+    model = model.to(device)
+    print("Model #Params: {:d}K".format(sum([x.nelement() for x in model.parameters()]) // 1000))
+
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+    scheduler = lr_scheduler.ExponentialLR(optimizer, args.gamma)
+
+    dur = []
+    t0 = time.time()
+    for epoch in range(args.max_epoch):
+        word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
+        for it, (batch_trees, batch_tree_graphs, batch_mol_graphs,
+                 stereo_cand_batch_idx, stereo_cand_labels, batch_stereo_cand_graphs) \
+                in enumerate(dataloader):
+            batch_tree_graphs = batch_tree_graphs.to(device)
+            batch_mol_graphs = batch_mol_graphs.to(device)
+            stereo_cand_batch_idx = stereo_cand_batch_idx.to(device)
+            batch_stereo_cand_graphs = batch_stereo_cand_graphs.to(device)
+
+            loss, kl_div, wacc, tacc, sacc, dacc = model(
+                batch_trees, batch_tree_graphs, batch_mol_graphs, stereo_cand_batch_idx,
+                stereo_cand_labels, batch_stereo_cand_graphs, beta=args.beta)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            word_acc += wacc
+            topo_acc += tacc
+            assm_acc += sacc
+            steo_acc += dacc
+
+            if (it + 1) % args.print_iter == 0:
+                dur.append(time.time() - t0)
+                word_acc = word_acc / args.print_iter * 100
+                topo_acc = topo_acc / args.print_iter * 100
+                assm_acc = assm_acc / args.print_iter * 100
+                steo_acc = steo_acc / args.print_iter * 100
+
+                print('Epoch {:d}/{:d} | Iter {:d}/{:d} | KL: {:.1f}, Word: {:.2f}, '
+                      'Topo: {:.2f}, Assm: {:.2f}, Steo: {:.2f} | '
+                      'Estimated time per epoch: {:.4f}'.format(
+                    epoch + 1, args.max_epoch, it + 1, len(dataloader), kl_div, word_acc,
+                    topo_acc, assm_acc, steo_acc, np.mean(dur) / args.print_iter * len(dataloader)))
+                word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
+                sys.stdout.flush()
+                t0 = time.time()
+
+            if (it + 1) % 15000 == 0:
+                scheduler.step()
+
+            if (it + 1) % 1000 == 0:
+                torch.save(model.state_dict(), args.save_path + "/model.iter-{:d}-{:d}".format(epoch, it + 1))
+
+        scheduler.step()
+        torch.save(model.state_dict(), args.save_path + "/model.iter-" + str(epoch))
+
+if __name__ == '__main__':
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser()
+    parser.add_argument('-tr', '--train-path', type=str,
+                        help='Path to the training molecules, with one SMILES string a line')
+    parser.add_argument('-s', '--save-path', type=str, default='vae_model',
+                        help='Directory to save model checkpoints')
+    parser.add_argument('-m', '--model-path', type=str,
+                        help='Path to pre-trained model checkpoint')
+    parser.add_argument('-b', '--batch-size', type=int, default=40,
+                        help='Batch size')
+    parser.add_argument('-w', '--hidden-size', type=int, default=450,
+                        help='Hidden size')
+    parser.add_argument('-l', '--latent-size', type=int, default=56,
+                        help='Latent size')
+    parser.add_argument('-d', '--depth', type=int, default=3,
+                        help='Number of GNN layers')
+    parser.add_argument('-z', '--beta', type=float, default=0.001,
+                        help='Weight for KL loss term')
+    parser.add_argument('-lr', '--lr', type=float, default=0.0007,
+                        help='Learning rate')
+    parser.add_argument('-g', '--gamma', type=float, default=0.9,
+                        help='Multiplicative factor for learning rate decay')
+    parser.add_argument('-me', '--max-epoch', type=int, default=7,
+                        help='Maximum number of epochs for training')
+    parser.add_argument('-nw', '--num-workers', type=int, default=4,
+                        help='Number of subprocesses for data loading')
+    parser.add_argument('-pi', '--print-iter', type=int, default=20,
+                        help='Frequency for printing evaluation metrics')
+    parser.add_argument('-cpu', '--use-cpu', action='store_true',
+                        help='By default, the script uses GPU whenever available. '
+                             'This flag enforces the use of CPU.')
+    args = parser.parse_args()
+
+    main(args)
diff --git a/examples/generative_models/jtvae/vocab.py b/examples/generative_models/jtvae/vocab.py
deleted file mode 100644
index ab307ea0..00000000
--- a/examples/generative_models/jtvae/vocab.py
+++ /dev/null
@@ -1,52 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Generate vocabulary for a new dataset.
-
-if __name__ == '__main__':
-    import argparse
-    import os
-    import rdkit
-
-    from dgl.data.utils import _get_dgl_url, download, get_download_dir, extract_archive
-
-    from dgllife.data.jtvae import DGLMolTree
-
-    parser = argparse.ArgumentParser('Generate vocabulary for a molecule dataset')
-    parser.add_argument('-d', '--data-path', type=str,
-                        help='Path to the dataset')
-    parser.add_argument('-v', '--vocab', type=str,
-                        help='Path to the vocabulary file to save')
-    args = parser.parse_args()
-
-    lg = rdkit.RDLogger.logger()
-    lg.setLevel(rdkit.RDLogger.CRITICAL)
-
-    vocab = set()
-    with open(args.data_path, 'r') as f:
-        for line in f:
-            smiles = line.strip()
-            mol = DGLMolTree(smiles)
-            for i in mol.nodes_dict:
-                vocab.add(mol.nodes_dict[i]['smiles'])
-
-    with open(args.vocab, 'w') as f:
-        for v in vocab:
-            f.write(v + '\n')
-
-    # Get the vocabulary used for the pre-trained model
-    default_dir = get_download_dir()
-    vocab_file = '{}/jtvae/{}.txt'.format(default_dir, 'vocab')
-    if not os.path.exists(vocab_file):
-        zip_file_path = '{}/jtvae.zip'.format(default_dir)
-        download(_get_dgl_url('dataset/jtvae.zip'), path=zip_file_path)
-        extract_archive(zip_file_path, '{}/jtvae'.format(default_dir))
-    default_vocab = set()
-    with open(vocab_file, 'r') as f:
-        for line in f:
-            default_vocab.add(line.strip())
-
-    print('The new vocabulary is a subset of the default vocabulary: {}'.format(
-        vocab.issubset(default_vocab)))
diff --git a/python/dgllife/data/jtvae.py b/python/dgllife/data/jtvae.py
index c3385c8b..47b4674d 100644
--- a/python/dgllife/data/jtvae.py
+++ b/python/dgllife/data/jtvae.py
@@ -3,790 +3,31 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
 # Dataset for JTVAE
 
-import numpy as np
+import dgl
+import dgl.function as fn
+import itertools
 import torch
 
-from collections import defaultdict
-from dgl import batch, graph
 from dgl.data.utils import get_download_dir, _get_dgl_url, download, extract_archive
 from functools import partial
 from rdkit import Chem
-from rdkit.Chem.EnumerateStereoisomers import EnumerateStereoisomers
-from scipy.sparse import csr_matrix
-from scipy.sparse.csgraph import minimum_spanning_tree
 from torch.utils.data import Dataset
 
-from ..utils.featurizers import BaseAtomFeaturizer, BaseBondFeaturizer, ConcatFeaturizer, \
-    atom_type_one_hot, atom_degree_one_hot, atom_formal_charge_one_hot, atom_chiral_tag_one_hot, \
-    atom_is_aromatic, bond_type_one_hot, bond_is_in_ring, bond_stereo_one_hot
+from ..utils.featurizers import BaseAtomFeaturizer, ConcatFeaturizer, atom_type_one_hot, \
+    atom_degree_one_hot, atom_formal_charge_one_hot, atom_chiral_tag_one_hot, atom_is_aromatic, \
+    BaseBondFeaturizer, bond_type_one_hot, bond_is_in_ring, bond_stereo_one_hot
+from ..utils.jtvae.chemutils import get_mol
+from ..utils.jtvae.mol_tree import MolTree
 from ..utils.mol_to_graph import mol_to_bigraph
 
-__all__ = ['DGLMolTree',
-           'JTVAEDataset',
+__all__ = ['JTVAEDataset',
+           'JTVAEZINC',
            'JTVAECollator']
 
-def get_mol(smiles, kekulize=True):
-    """Convert the SMILES string into an RDKit molecule object
-
-    Parameters
-    ----------
-    smiles : str
-        A SMILES string.
-    kekulize : bool
-        Whether to kekulize the molecule.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Mol or None
-        A Kekulized RDKit molecule object if the input SMILES string is valid and None otherwise.
-    """
-    mol = Chem.MolFromSmiles(smiles)
-    if mol is None:
-        return None
-    if kekulize:
-        Chem.Kekulize(mol)
-    return mol
-
-def get_smiles(mol):
-    """Convert an RDKit molecule object into a SMILES string.
-
-    By default, the molecule is kekulized.
-
-    Parameters
-    ----------
-    mol : rdkit.Chem.rdchem.Mol
-        A molecule.
-
-    Returns
-    -------
-    str
-        The SMILES string corresponding to the molecule.
-    """
-    return Chem.MolToSmiles(mol, kekuleSmiles=True)
-
-def decode_stereo(smiles_2d):
-    """Get possible 3D SMILES by enumerating stereoisomers
-
-    Parameters
-    ----------
-    smiles_2d : str
-        A 2D SMILES string.
-
-    Returns
-    -------
-    list of str
-        List of possible 3D SMILES strings.
-    """
-    mol = Chem.MolFromSmiles(smiles_2d)
-    # Get all possible stereoisomers for a molecule
-    dec_isomers = list(EnumerateStereoisomers(mol))
-    dec_isomers = [Chem.MolFromSmiles(Chem.MolToSmiles(
-        mol, isomericSmiles=True)) for mol in dec_isomers]
-    smiles_3d = [Chem.MolToSmiles(mol, isomericSmiles=True)
-                for mol in dec_isomers]
-
-    chiral_n = [atom.GetIdx() for atom in dec_isomers[0].GetAtoms()
-                if int(atom.GetChiralTag()) > 0 and atom.GetSymbol() == "N"]
-
-    if len(chiral_n) > 0:
-        for mol in dec_isomers:
-            for idx in chiral_n:
-                mol.GetAtomWithIdx(idx).SetChiralTag(
-                    Chem.rdchem.ChiralType.CHI_UNSPECIFIED)
-            smiles_3d.append(Chem.MolToSmiles(mol, isomericSmiles=True))
-
-    return smiles_3d
-
-def set_atommap(mol, num=0):
-    """Set the atom map number for all atoms in the molecule.
-
-    Parameters
-    ----------
-    mol : rdkit.Chem.rdchem.Mol
-        A molecule.
-    num : int
-        The atom map number to set for all atoms. If 0, it will
-        clear the atom map.
-    """
-    for atom in mol.GetAtoms():
-        atom.SetAtomMapNum(num)
-
-def sanitize(mol):
-    """Sanitize the molecule.
-
-    Parameters
-    ----------
-    mol : rdkit.Chem.rdchem.Mol
-        A molecule.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Mol
-        Sanitized molecule.
-    """
-    try:
-        smiles = get_smiles(mol)
-        mol = get_mol(smiles)
-    except Exception:
-        return None
-    return mol
-
-def copy_atom(atom):
-    """Get a deep copy of an atom.
-
-    Parameters
-    ----------
-    atom : rdkit.Chem.rdchem.Atom
-        An atom.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Atom
-        A deep copy of the input atom object.
-    """
-    new_atom = Chem.Atom(atom.GetSymbol())
-    new_atom.SetFormalCharge(atom.GetFormalCharge())
-    new_atom.SetAtomMapNum(atom.GetAtomMapNum())
-    return new_atom
-
-def copy_edit_mol(mol):
-    """Get a deep copy of a molecule.
-
-    Parameters
-    ----------
-    mol : rdkit.Chem.rdchem.Mol
-        A molecule.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Mol
-        A deep copy of the input molecule object.
-    """
-    new_mol = Chem.RWMol(Chem.MolFromSmiles(''))
-    for atom in mol.GetAtoms():
-        new_atom = copy_atom(atom)
-        new_mol.AddAtom(new_atom)
-    for bond in mol.GetBonds():
-        a1 = bond.GetBeginAtom().GetIdx()
-        a2 = bond.GetEndAtom().GetIdx()
-        bt = bond.GetBondType()
-        new_mol.AddBond(a1, a2, bt)
-    return new_mol
-
-def get_clique_mol(mol, atoms):
-    """Get an RDKit molecule object for a fragment of the molecule.
-
-    Parameters
-    ----------
-    mol : rdkit.Chem.rdchem.Mol
-        A molecule.
-    atoms : list of int
-        List of ids for the atoms in the fragment.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Mol
-        The RDKit molecule object corresponding to the fragment.
-    """
-    smiles = Chem.MolFragmentToSmiles(mol, atoms, kekuleSmiles=True)
-    new_mol = Chem.MolFromSmiles(smiles, sanitize=False)
-    new_mol = copy_edit_mol(new_mol).GetMol()
-    new_mol = sanitize(new_mol)  # We assume this is not None
-    return new_mol
-
-def get_atom_to_substructures(cliques):
-    """Get the substructures that each atom belongs to.
-
-    Parameters
-    ----------
-    cliques : list
-        Each element is a list of int, representing a
-        non-ring bond or a simple cycle.
-
-    Returns
-    -------
-    nei_list : dict
-        Mapping atom ids to the ids of the substructures that the atom belongs to.
-    """
-    nei_list = defaultdict(list)
-    for i in range(len(cliques)):
-        for atom in cliques[i]:
-            nei_list[atom].append(i)
-
-    return nei_list
-
-def tree_decomp(mol, mst_max_weight=100):
-    """Tree decomposition of a molecule for junction tree construction.
-
-    Parameters
-    ----------
-    mol : rdkit.Chem.rdchem.Mol
-        A molecule.
-    mst_max_weight : int
-        Max weight considered in generating a minimum spanning tree
-
-    Returns
-    -------
-    list
-        Clusters. Each element is a list of int,
-        representing the atoms that constitute the cluster.
-    list
-        Edges between the clusters. Each element is a 2-tuple of cluster IDs.
-    """
-    n_atoms = mol.GetNumAtoms()
-    if n_atoms == 1:
-        return [[0]], []
-
-    cliques = []
-
-    # Find all edges not belonging to any cycles
-    for bond in mol.GetBonds():
-        a1 = bond.GetBeginAtom().GetIdx()
-        a2 = bond.GetEndAtom().GetIdx()
-        if not bond.IsInRing():
-            cliques.append([a1, a2])
-
-    # Find all simple cycles, each represented by a list of IDs of the atoms in the ring
-    ssr = [list(x) for x in Chem.GetSymmSSSR(mol)]
-    cliques.extend(ssr)
-
-    # Record the non-ring bonds/simple cycles that each atom belongs to
-    nei_list = get_atom_to_substructures(cliques)
-
-    # Merge simple rings that have more than two overlapping atoms
-    for i in range(len(cliques)):
-        if len(cliques[i]) <= 2:
-            continue
-        for atom in cliques[i]:
-            for j in nei_list[atom]:
-                if i >= j or len(cliques[j]) <= 2:
-                    continue
-                inter = set(cliques[i]) & set(cliques[j])
-                if len(inter) > 2:
-                    cliques[i].extend(cliques[j])
-                    cliques[i] = list(set(cliques[i]))
-                    cliques[j] = []
-
-    # Remove merged simple cycles
-    cliques = [c for c in cliques if len(c) > 0]
-    # Record the non-ring bonds/simple cycles that each atom belongs to
-    nei_list = get_atom_to_substructures(cliques)
-
-    # Build edges and add singleton cliques
-    edges = defaultdict(int)
-    for atom in range(n_atoms):
-        if len(nei_list[atom]) <= 1:
-            continue
-        cnei = nei_list[atom]
-        bonds = [c for c in cnei if len(cliques[c]) == 2]
-        rings = [c for c in cnei if len(cliques[c]) > 4]
-        # In general, if len(cnei) >= 3, a singleton should be added,
-        # but 1 bond + 2 ring is currently not dealt with.
-        if len(bonds) > 2 or (len(bonds) == 2 and len(cnei) > 2):
-            # Add singleton clique
-            cliques.append([atom])
-            c2 = len(cliques) - 1
-            for c1 in cnei:
-                edges[(c1, c2)] = 1
-        elif len(rings) > 2:  # Multiple (n>2) complex rings
-            cliques.append([atom])
-            c2 = len(cliques) - 1
-            for c1 in cnei:
-                edges[(c1, c2)] = mst_max_weight - 1
-        else:
-            for i in range(len(cnei)):
-                for j in range(i + 1, len(cnei)):
-                    c1, c2 = cnei[i], cnei[j]
-                    inter = set(cliques[c1]) & set(cliques[c2])
-                    if edges[(c1, c2)] < len(inter):
-                        # cnei[i] < cnei[j] by construction
-                        edges[(c1, c2)] = len(inter)
-
-    edges = [u + (mst_max_weight-v,) for u, v in edges.items()]
-    if len(edges) == 0:
-        return cliques, edges
-
-    # Compute Maximum Spanning Tree
-    row, col, data = list(zip(*edges))
-    n_clique = len(cliques)
-    clique_graph = csr_matrix((data, (row, col)), shape=(n_clique, n_clique))
-    junc_tree = minimum_spanning_tree(clique_graph)
-    row, col = junc_tree.nonzero()
-    edges = [(row[i], col[i]) for i in range(len(row))]
-
-    return (cliques, edges)
-
-def atom_equal(a1, a2):
-    """Check if two atoms have the same type and formal charge.
-
-    Parameters
-    ----------
-    a1 : rdkit.Chem.rdchem.Atom
-        An RDKit atom object.
-    a2 : rdkit.Chem.rdchem.Atom
-        An RDKit atom object.
-
-    Returns
-    -------
-    bool
-        Whether they are equal.
-    """
-    return a1.GetSymbol() == a2.GetSymbol() and a1.GetFormalCharge() == a2.GetFormalCharge()
-
-def ring_bond_equal(b1, b2, reverse=False):
-    """Check whether the end atoms of the two bonds have the same type and formal charge.
-
-    Parameters
-    ----------
-    b1 : rdkit.Chem.rdchem.Bond
-        An RDKit bond object.
-    b2 : rdkit.Chem.rdchem.Bond
-        An RDKit bond object.
-    reverse : bool
-        Whether to reverse b2 in comparison.
-
-    Returns
-    -------
-    bool
-        Whether the bonds are equal.
-    """
-    b1 = (b1.GetBeginAtom(), b1.GetEndAtom())
-    if reverse:
-        b2 = (b2.GetEndAtom(), b2.GetBeginAtom())
-    else:
-        b2 = (b2.GetBeginAtom(), b2.GetEndAtom())
-    return atom_equal(b1[0], b2[0]) and atom_equal(b1[1], b2[1])
-
-def attach_mols_nx(ctr_mol, neighbors, prev_nodes, nei_amap):
-    """Attach clusters to a central molecule.
-
-    Parameters
-    ----------
-    ctr_mol : rdkit.Chem.rdchem.Mol
-        The central molecule.
-    neighbors : list of dict
-        Each element contains the information of a neighboring cluster.
-    prev_nodes : list of dict
-        Each element contains the information of a previous cluster.
-    nei_amap : dict
-        nei_amap[nei_id][nei_atom] maps an atom in a neighboring cluster
-        with id nei_id to an atom in the central molecule.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Mol
-        The central molecule with clusters attached.
-    """
-    prev_nids = [node['nid'] for node in prev_nodes]
-    for nei_node in prev_nodes + neighbors:
-        nei_id, nei_mol = nei_node['nid'], nei_node['mol']
-        amap = nei_amap[nei_id]
-        for atom in nei_mol.GetAtoms():
-            if atom.GetIdx() not in amap:
-                new_atom = copy_atom(atom)
-                amap[atom.GetIdx()] = ctr_mol.AddAtom(new_atom)
-
-        if nei_mol.GetNumBonds() == 0:
-            nei_atom = nei_mol.GetAtomWithIdx(0)
-            ctr_atom = ctr_mol.GetAtomWithIdx(amap[0])
-            ctr_atom.SetAtomMapNum(nei_atom.GetAtomMapNum())
-        else:
-            for bond in nei_mol.GetBonds():
-                a1 = amap[bond.GetBeginAtom().GetIdx()]
-                a2 = amap[bond.GetEndAtom().GetIdx()]
-                if ctr_mol.GetBondBetweenAtoms(a1, a2) is None:
-                    ctr_mol.AddBond(a1, a2, bond.GetBondType())
-                elif nei_id in prev_nids:  # father node overrides
-                    ctr_mol.RemoveBond(a1, a2)
-                    ctr_mol.AddBond(a1, a2, bond.GetBondType())
-    return ctr_mol
-
-def local_attach_nx(ctr_mol, neighbors, prev_nodes, amap_list):
-    """Pre-process and attach clusters to a central molecule.
-
-    Parameters
-    ----------
-    ctr_mol : rdkit.Chem.rdchem.Mol
-        The central molecule.
-    neighbors : list of dict
-        Each element contains the information of a neighboring cluster.
-    prev_nodes : list of dict
-        Each element contains the information of a neighboring cluster.
-    amap_list : list of 3-tuples
-        Each tuple consists of the id of the neighboring cluster,
-        the id of the atom in the central molecule and the id of the same atom in the
-        neighboring cluster.
-
-    Returns
-    -------
-    rdkit.Chem.rdchem.Mol
-        The central molecule with the clusters attached.
-    """
-    ctr_mol = copy_edit_mol(ctr_mol)
-    nei_amap = {nei['nid']: {} for nei in prev_nodes + neighbors}
-
-    for nei_id, ctr_atom, nei_atom in amap_list:
-        nei_amap[nei_id][nei_atom] = ctr_atom
-
-    ctr_mol = attach_mols_nx(ctr_mol, neighbors, prev_nodes, nei_amap)
-    return ctr_mol.GetMol()
-
-def enum_attach_nx(ctr_mol, nei_node, amap, singletons):
-    """Enumerate possible ways to attach a cluster to the central molecule.
-
-    Parameters
-    ----------
-    ctr_mol : rdkit.Chem.rdchem.Mol
-        The central molecule.
-    nei_node : dict
-        Information for a cluster to attach to the central molecule.
-    amap : list of 3-tuples
-        Each tuple consists of the id of the neighboring cluster,
-        the id of the atom in the central molecule and the id of the same atom in the
-        neighboring cluster.
-    singletons : list of int
-        IDs for the neighboring singletons attached.
-
-    Returns
-    -------
-    list
-        Each element is of the form "amap", corresponding to an attachment configuration.
-    """
-    nei_mol, nei_idx = nei_node['mol'], nei_node['nid']
-    att_confs = []
-    # A black list of atoms in the central molecule connected to singletons
-    black_list = [atom_idx for nei_id, atom_idx,
-                  _ in amap if nei_id in singletons]
-    ctr_atoms = [atom for atom in ctr_mol.GetAtoms() if atom.GetIdx()
-                 not in black_list]
-    ctr_bonds = [bond for bond in ctr_mol.GetBonds()]
-
-    if nei_mol.GetNumBonds() == 0:  # neighbor singleton
-        nei_atom = nei_mol.GetAtomWithIdx(0)
-        used_list = [atom_idx for _, atom_idx, _ in amap]
-        for atom in ctr_atoms:
-            if atom_equal(atom, nei_atom) and atom.GetIdx() not in used_list:
-                new_amap = amap + [(nei_idx, atom.GetIdx(), 0)]
-                att_confs.append(new_amap)
-
-    elif nei_mol.GetNumBonds() == 1:  # neighbor is a bond
-        bond = nei_mol.GetBondWithIdx(0)
-        bond_val = int(bond.GetBondTypeAsDouble())
-        b1, b2 = bond.GetBeginAtom(), bond.GetEndAtom()
-
-        for atom in ctr_atoms:
-            # Optimize if atom is carbon (other atoms may change valence)
-            if atom.GetAtomicNum() == 6 and atom.GetTotalNumHs() < bond_val:
-                continue
-            if atom_equal(atom, b1):
-                new_amap = amap + [(nei_idx, atom.GetIdx(), b1.GetIdx())]
-                att_confs.append(new_amap)
-            elif atom_equal(atom, b2):
-                new_amap = amap + [(nei_idx, atom.GetIdx(), b2.GetIdx())]
-                att_confs.append(new_amap)
-    else:
-        # intersection is an atom
-        for a1 in ctr_atoms:
-            for a2 in nei_mol.GetAtoms():
-                if atom_equal(a1, a2):
-                    # Optimize if atom is carbon (other atoms may change valence)
-                    if a1.GetAtomicNum() == 6 and a1.GetTotalNumHs() + a2.GetTotalNumHs() < 4:
-                        continue
-                    new_amap = amap + [(nei_idx, a1.GetIdx(), a2.GetIdx())]
-                    att_confs.append(new_amap)
-
-        # intersection is an bond
-        if ctr_mol.GetNumBonds() > 1:
-            for b1 in ctr_bonds:
-                for b2 in nei_mol.GetBonds():
-                    if ring_bond_equal(b1, b2):
-                        new_amap = amap + [(nei_idx,
-                                            b1.GetBeginAtom().GetIdx(),
-                                            b2.GetBeginAtom().GetIdx()),
-                                           (nei_idx,
-                                            b1.GetEndAtom().GetIdx(),
-                                            b2.GetEndAtom().GetIdx())]
-                        att_confs.append(new_amap)
-
-                    if ring_bond_equal(b1, b2, reverse=True):
-                        new_amap = amap + [(nei_idx,
-                                            b1.GetBeginAtom().GetIdx(),
-                                            b2.GetEndAtom().GetIdx()),
-                                           (nei_idx,
-                                            b1.GetEndAtom().GetIdx(),
-                                            b2.GetBeginAtom().GetIdx())]
-                        att_confs.append(new_amap)
-    return att_confs
-
-def enum_assemble_nx(node, neighbors, prev_nodes=None, prev_amap=None):
-    """Enumerate all possible ways to assemble the central molecule and the neighboring clusters.
-
-    Parameters
-    ----------
-    node : dict
-        The information of the central molecule.
-    neighbors : list of dict
-        Each element contains the information of a neighboring cluster.
-    prev_nodes : list of dict
-        Each element contains the information of a neighboring cluster.
-    prev_amap : list of 3-tuples
-        Each tuple consists of the id of the neighboring cluster,
-        the id of the atom in the central molecule and the id of the same atom in the
-        neighboring cluster.
-
-    Returns
-    -------
-    list of 3-tuples
-        Each 3-tuple consists of the SMILES, molecule and amap corresponding to a molecule
-        assembled from the central molecule and the neighboring cluster. The amap is a list
-        of tuples. Each tuple consists of the id of the neighboring cluster,
-        the id of the atom in the central molecule and the id of the same atom in the
-        neighboring cluster.
-    """
-    if prev_nodes is None:
-        prev_nodes = []
-    if prev_amap is None:
-        prev_amap = []
-    all_attach_confs = []
-    singletons = [nei_node['nid'] for nei_node in neighbors +
-                  prev_nodes if nei_node['mol'].GetNumAtoms() == 1]
-
-    def search(cur_amap, depth, max_ncand=2000):
-        if len(all_attach_confs) > max_ncand:
-            return
-        if depth == len(neighbors):
-            all_attach_confs.append(cur_amap)
-            return
-
-        nei_node = neighbors[depth]
-        # Enumerate possible ways to attach nei_node to the central molecule
-        cand_amap = enum_attach_nx(node['mol'], nei_node, cur_amap, singletons)
-        cand_smiles = set()
-        candidates = []
-        for amap in cand_amap:
-            cand_mol = local_attach_nx(
-                node['mol'], neighbors[:depth+1], prev_nodes, amap)
-            cand_mol = sanitize(cand_mol)
-            if cand_mol is None:
-                continue
-            smiles = Chem.MolToSmiles(cand_mol, kekuleSmiles=True)
-            if smiles in cand_smiles:
-                continue
-            cand_smiles.add(smiles)
-            candidates.append(amap)
-
-        if len(candidates) == 0:
-            return []
-
-        for new_amap in candidates:
-            search(new_amap, depth + 1)
-
-    search(prev_amap, 0)
-    cand_smiles = set()
-    candidates = []
-    for amap in all_attach_confs:
-        cand_mol = local_attach_nx(node['mol'], neighbors, prev_nodes, amap)
-        cand_mol = Chem.MolFromSmiles(Chem.MolToSmiles(cand_mol))
-        smiles = Chem.MolToSmiles(cand_mol)
-        if smiles in cand_smiles:
-            continue
-        cand_smiles.add(smiles)
-        Chem.Kekulize(cand_mol)
-        candidates.append((smiles, cand_mol, amap))
-
-    return candidates
-
-class DGLMolTree():
-    """Junction tree.
-
-    Parameters
-    ----------
-    smiles : str
-        A SMILES string.
-    """
-    def __init__(self, smiles):
-        self.smiles = smiles
-        self.nodes_dict = {}
-
-        if smiles is None:
-            self.g = graph(([], []), idtype=torch.int32)
-            return
-
-        self.mol = get_mol(smiles)
-
-        # Stereo Generation
-        mol = Chem.MolFromSmiles(smiles)
-        self.smiles_3d = Chem.MolToSmiles(mol, isomericSmiles=True)
-        self.smiles_2d = Chem.MolToSmiles(mol)
-        self.stereo_cands = decode_stereo(self.smiles_2d)
-
-        # Junction tree construction
-        cliques, edges = tree_decomp(self.mol)
-        root = 0
-        for i, c in enumerate(cliques):
-            cmol = get_clique_mol(self.mol, c)
-            csmiles = Chem.MolToSmiles(cmol, kekuleSmiles=True)
-            self.nodes_dict[i] = dict(
-                smiles=csmiles,
-                mol=get_mol(csmiles),
-                clique=c,
-            )
-            if min(c) == 0:
-                root = i
-
-        # The clique with atom ID 0 becomes root
-        if root > 0:
-            for attr in self.nodes_dict[0]:
-                self.nodes_dict[0][attr], self.nodes_dict[root][attr] = \
-                    self.nodes_dict[root][attr], self.nodes_dict[0][attr]
-
-        src = np.zeros((len(edges) * 2,), dtype='int')
-        dst = np.zeros((len(edges) * 2,), dtype='int')
-        for i, (_x, _y) in enumerate(edges):
-            if _x == root:
-                x = 0
-            elif _x == 0:
-                x = root
-            else:
-                x = _x
-
-            if _y == root:
-                y = 0
-            elif _y == 0:
-                y = root
-            else:
-                y = _y
-
-            src[2 * i] = x
-            dst[2 * i] = y
-            src[2 * i + 1] = y
-            dst[2 * i + 1] = x
-
-        self.g = graph((src, dst), num_nodes=len(cliques), idtype=torch.int32)
-
-        for i in self.nodes_dict:
-            self.nodes_dict[i]['nid'] = i + 1
-            if self.g.out_degrees(i) > 1:  # Leaf node mol is not marked
-                set_atommap(self.nodes_dict[i]['mol'], self.nodes_dict[i]['nid'])
-            self.nodes_dict[i]['is_leaf'] = (self.g.out_degrees(i) == 1)
-
-    def treesize(self):
-        """Get the number of nodes in the junction tree.
-
-        Returns
-        -------
-        int
-            Get the number of nodes (clusters) in the junction tree.
-        """
-        return self.g.num_nodes()
-
-    def _recover_node(self, i, original_mol):
-        """Get the SMILES string corresponding to the i-th cluster in the
-        original molecule.
-
-        Parameters
-        ----------
-        i : int
-            The id of a cluster.
-        original_mol : rdkit.Chem.rdchem.Mol
-            The original molecule.
-
-        Returns
-        -------
-        str
-            A SMILES string.
-        """
-        node = self.nodes_dict[i]
-
-        clique = []
-        clique.extend(node['clique'])
-        if not node['is_leaf']:
-            for cidx in node['clique']:
-                original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(node['nid'])
-
-        for j in self.g.successors(i).numpy():
-            nei_node = self.nodes_dict[j]
-            clique.extend(nei_node['clique'])
-            if nei_node['is_leaf']:  # Leaf node, no need to mark
-                continue
-            for cidx in nei_node['clique']:
-                # allow singleton node override the atom mapping
-                if cidx not in node['clique'] or len(nei_node['clique']) == 1:
-                    atom = original_mol.GetAtomWithIdx(cidx)
-                    atom.SetAtomMapNum(nei_node['nid'])
-
-        clique = list(set(clique))
-        label_mol = get_clique_mol(original_mol, clique)
-        node['label'] = Chem.MolToSmiles(Chem.MolFromSmiles(
-            Chem.MolToSmiles(label_mol, kekuleSmiles=True)))
-        node['label_mol'] = get_mol(node['label'])
-
-        for cidx in clique:
-            original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(0)
-
-        return node['label']
-
-    def recover(self):
-        """Get the SMILES strings corresponding to all clusters in the
-        original molecule."""
-        for i in self.nodes_dict:
-            self._recover_node(i, self.mol)
-
-    def _assemble_node(self, i):
-        """Assemble a cluster with its successors.
-
-        Parameters
-        ----------
-        i : int
-            The id of a cluster.
-        """
-        neighbors = [self.nodes_dict[j] for j in self.g.successors(i).numpy()
-                     if self.nodes_dict[j]['mol'].GetNumAtoms() > 1]
-        neighbors = sorted(neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
-        singletons = [self.nodes_dict[j] for j in self.g.successors(i).numpy()
-                      if self.nodes_dict[j]['mol'].GetNumAtoms() == 1]
-        # All successors sorted based on their size
-        neighbors = singletons + neighbors
-
-        cands = enum_assemble_nx(self.nodes_dict[i], neighbors)
-
-        if len(cands) > 0:
-            self.nodes_dict[i]['cands'], self.nodes_dict[i]['cand_mols'], _ = list(zip(*cands))
-            self.nodes_dict[i]['cands'] = list(self.nodes_dict[i]['cands'])
-            self.nodes_dict[i]['cand_mols'] = list(self.nodes_dict[i]['cand_mols'])
-        else:
-            self.nodes_dict[i]['cands'] = []
-            self.nodes_dict[i]['cand_mols'] = []
-
-    def assemble(self):
-        """Assemble each cluster with its successors"""
-        for i in self.nodes_dict:
-            self._assemble_node(i)
-
-def _set_node_id(mol_tree, vocab):
-    """Get the id corresponding to each cluster in the vocabulary
-
-    Parameters
-    ----------
-    mol_tree : MolTree
-        A junction tree for a molecule.
-    vocab : Vocab
-        A vocabulary.
-
-    Returns
-    -------
-    list of int
-        The i-th element corresponds to the id of the i-th cluster in the vocabulary.
-    """
-    wid = []
-    for i, node in enumerate(mol_tree.nodes_dict):
-        mol_tree.nodes_dict[node]['idx'] = i
-        wid.append(vocab.get_index(mol_tree.nodes_dict[node]['smiles']))
-
-    return wid
-
 def get_atom_featurizer_enc():
     """Get the atom featurizer for encoding.
 
@@ -866,440 +107,262 @@ def get_bond_featurizer_dec():
     ])})
     return featurizer
 
-def mol2dgl_enc(smiles, atom_featurizer, bond_featurizer):
-    """Convert a SMILES to a DGLGraph for encoding.
-
-    Parameters
-    ----------
-    smiles : str
-        A SMILES string.
-    atom_featurizer : callable
-        Function for featurizing atoms.
-    bond_featurizer : callable
-        Function for featurizing bonds.
-
-    Returns
-    -------
-    g : DGLGraph
-        The DGLGraph for encoding. g.ndata['x'] stores the atom features and
-        g.edata['x'] stores the bond features.
-    """
-    mol = get_mol(smiles)
-    return mol_to_bigraph(mol=mol,
-                          node_featurizer=atom_featurizer,
-                          edge_featurizer=bond_featurizer,
-                          canonical_atom_order=False)
-
-def mol2dgl_dec(cand_batch, atom_featurizer, bond_featurizer):
-    """Convert a batch of candidate molecules to DGLGraphs for decoding.
-
-    Parameters
-    ----------
-    cand_batch : list of 3-tuples
-        Each tuple consists of the candidate molecule, the junction tree of
-        the original molecule and the id for the cluster.
-    atom_featurizer : callable
-        Function for featurizing atoms.
-    bond_featurizer : callable
-        Function for featurizing bonds.
-
-    Returns
-    -------
-    g_list : list of DGLGraph
-        Each graph corresponds to a candidate molecule decoded. g_list[i].ndata['x']
-        stores the atom features and g_list[i].edata['x'] stores the bond features.
-    1D LongTensor
-        The edges in the junction tree corresponding to the edges in ``g_list``.
-    1D LongTensor
-        The edges in the batched ``g_list``.
-    1D LongTensor
-        The nodes in the batched ``g_list``.
-    """
-    # Note that during graph decoding they don't predict stereochemistry-related
-    # characteristics (i.e. Chiral Atoms, E-Z, Cis-Trans).  Instead, they decode
-    # the 2-D graph first, then enumerate all possible 3-D forms and find the
-    # one with highest score.
-    cand_graphs = []
-    tree_mess_source_edges = []  # map these edges from trees to...
-    tree_mess_target_edges = []  # these edges on candidate graphs
-    tree_mess_target_nodes = []
-    n_nodes = 0
-
-    for mol, mol_tree, ctr_node_id in cand_batch:
-        n_atoms = mol.GetNumAtoms()
-
-        g = mol_to_bigraph(mol,
-                           node_featurizer=atom_featurizer,
-                           edge_featurizer=bond_featurizer,
-                           canonical_atom_order=False)
-        cand_graphs.append(g)
-
-        if isinstance(mol_tree, DGLMolTree):
-            tree_graph = mol_tree.g
-        else:
-            tree_graph = mol_tree
-
-        for i, bond in enumerate(mol.GetBonds()):
-            a1, a2 = bond.GetBeginAtom(), bond.GetEndAtom()
-            begin_idx, end_idx = a1.GetIdx(), a2.GetIdx()
-            x_nid, y_nid = a1.GetAtomMapNum(), a2.GetAtomMapNum()
-            # Tree node ID in the batch
-            x_bid = mol_tree.nodes_dict[x_nid - 1]['idx'] if x_nid > 0 else -1
-            y_bid = mol_tree.nodes_dict[y_nid - 1]['idx'] if y_nid > 0 else -1
-
-            if x_bid >= 0 and y_bid >= 0 and x_bid != y_bid:
-                if tree_graph.has_edges_between(x_bid, y_bid):
-                    tree_mess_target_edges.append(
-                        (begin_idx + n_nodes, end_idx + n_nodes))
-                    tree_mess_source_edges.append((x_bid, y_bid))
-                    tree_mess_target_nodes.append(end_idx + n_nodes)
-                if tree_graph.has_edges_between(y_bid, x_bid):
-                    tree_mess_target_edges.append(
-                        (end_idx + n_nodes, begin_idx + n_nodes))
-                    tree_mess_source_edges.append((y_bid, x_bid))
-                    tree_mess_target_nodes.append(begin_idx + n_nodes)
-
-        # Update offset
-        n_nodes += n_atoms
-
-    return cand_graphs, \
-           torch.IntTensor(tree_mess_source_edges), \
-           torch.IntTensor(tree_mess_target_edges), \
-           torch.IntTensor(tree_mess_target_nodes)
-
 class JTVAEDataset(Dataset):
     """Dataset for JTVAE
 
-    JTVAE is introduced in `Junction Tree Variational Autoencoder for Molecular Graph Generation
-    <https://arxiv.org/abs/1802.04364>`__.
-
     Parameters
     ----------
-    data : str
-        * If 'train' or 'test', it will use the training or test subset of the ZINC dataset.
-        * Otherwise, it should be the path to a .txt file for a dataset. The .txt file should
-          contain one SMILES per line.
-    vocab : Vocab
-        Loaded vocabulary.
+    data_file : str
+        Path to a file of SMILES strings, with one SMILES string a line.
+    vocab : JTVAEVocab
+        Vocabulary for JTVAE.
+    cache : bool
+        Whether to cache the trees to speed up data loading or always construct trees on the fly.
     training : bool
         Whether the dataset is for training or not.
     """
-    def __init__(self, data, vocab, training=True):
-        dir = get_download_dir()
-
-        _url = _get_dgl_url('dataset/jtnn.zip')
-        zip_file_path = '{}/jtnn.zip'.format(dir)
-        download(_url, path=zip_file_path)
-        extract_archive(zip_file_path, '{}/jtnn'.format(dir))
-
-        print('Loading data...')
-        if data in ['train', 'test']:
-            # ZINC subset
-            data_file = '{}/jtnn/{}.txt'.format(dir, data)
-        else:
-            # New dataset
-            data_file = data
-        with open(data_file) as f:
+    def __init__(self, data_file, vocab, cache=False, training=True):
+        with open(data_file, 'r') as f:
             self.data = [line.strip("\r\n ").split()[0] for line in f]
         self.vocab = vocab
+        self.cache = cache
+        if cache:
+            self.trees = [None for _ in range(len(self))]
+            self.mol_graphs = [None for _ in range(len(self))]
+            self.stereo_cand_graphs = [None for _ in range(len(self))]
 
-        print('Loading finished')
-        print('\t# samples:', len(self.data))
         self.training = training
-
         self.atom_featurizer_enc = get_atom_featurizer_enc()
         self.bond_featurizer_enc = get_bond_featurizer_enc()
-        self.atom_featurizer_dec = get_atom_featurizer_dec()
-        self.bond_featurizer_dec = get_bond_featurizer_dec()
-
-    @staticmethod
-    def move_to_device(mol_batch, device):
-        """Move a data batch to the target device.
-
-        Parameters
-        ----------
-        mol_batch : dict
-            A batch of datapoints.
-        device
-            A target device.
-
-        Returns
-        -------
-        dict
-            The batch of datapoints moved to the target device.
-        """
-        trees = []
-        for tr in mol_batch['mol_trees']:
-            tr.g = tr.g.to(device)
-            trees.append(tr)
-        mol_batch['mol_trees'] = trees
-        mol_batch['mol_graph_batch'] = mol_batch['mol_graph_batch'].to(device)
-        if 'cand_graph_batch' in mol_batch:
-            mol_batch['cand_graph_batch'] = mol_batch['cand_graph_batch'].to(device)
-        if mol_batch.get('stereo_cand_graph_batch') is not None:
-            mol_batch['stereo_cand_graph_batch'] = mol_batch['stereo_cand_graph_batch'].to(device)
-
-        return mol_batch
 
     def __len__(self):
-        """Get the size of the dataset.
+        """Get the size of the dataset
 
         Returns
         -------
         int
-            Size of the dataset.
+            Number of molecules in the dataset.
         """
         return len(self.data)
 
+    def getitem_train(self, idx):
+        if self.cache and self.trees[idx] is not None:
+            mol_tree = self.trees[idx]
+            mol_graph = self.mol_graphs[idx]
+            stereo_cand_graphs = self.stereo_cand_graphs[idx]
+        else:
+            smiles = self.data[idx]
+            mol_tree = MolTree(smiles)
+            mol_tree.recover()
+            mol_tree.assemble()
+
+            for node_id, node in mol_tree.nodes_dict.items():
+                if node['label'] not in node['cands']:
+                    node['cands'].append(node['label'])
+                    node['cand_mols'].append(node['label_mol'])
+
+            wid = [self.vocab.get_index(mol_tree.nodes_dict[i]['smiles'])
+                   for i in mol_tree.nodes_dict]
+            mol_tree.graph.ndata['wid'] = torch.LongTensor(wid)
+
+            # Construct molecular graphs
+            mol = get_mol(smiles)
+            mol_graph = mol_to_bigraph(mol,
+                                       node_featurizer=self.atom_featurizer_enc,
+                                       edge_featurizer=self.bond_featurizer_enc,
+                                       canonical_atom_order=False)
+            mol_graph.apply_edges(fn.copy_u('x', 'src'))
+            mol_graph.edata['x'] = torch.cat(
+                [mol_graph.edata.pop('src'), mol_graph.edata['x']], dim=1)
+
+            stereo_cand_graphs = []
+            stereo_cands = mol_tree.stereo_cands
+            if len(stereo_cands) != 1:
+                if mol_tree.smiles3D not in stereo_cands:
+                    stereo_cands.append(mol_tree.smiles3D)
+
+                for cand in stereo_cands:
+                    cand = get_mol(cand)
+                    cg = mol_to_bigraph(cand, node_featurizer=self.atom_featurizer_enc,
+                                        edge_featurizer=self.bond_featurizer_enc,
+                                        canonical_atom_order=False)
+                    cg.apply_edges(fn.copy_u('x', 'src'))
+                    cg.edata['x'] = torch.cat([cg.edata.pop('src'), cg.edata['x']], dim=1)
+                    stereo_cand_graphs.append(cg)
+
+            if self.cache:
+                self.trees[idx] = mol_tree
+                self.mol_graphs[idx] = mol_graph
+                self.stereo_cand_graphs[idx] = stereo_cand_graphs
+
+        return mol_tree, mol_graph, stereo_cand_graphs
+
+    def getitem_non_train(self, idx):
+        if self.cache and self.trees[idx] is not None:
+            mol_tree = self.trees[idx]
+            mol_graph = self.mol_graphs[idx]
+        else:
+            smiles = self.data[idx]
+            mol = Chem.MolFromSmiles(smiles)
+            smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
+            mol_tree = MolTree(smiles)
+            mol_tree.recover()
+
+            wid = [self.vocab.get_index(mol_tree.nodes_dict[i]['smiles'])
+                   for i in mol_tree.nodes_dict]
+            mol_tree.graph.ndata['wid'] = torch.LongTensor(wid)
+
+            # Construct molecular graphs
+            mol = get_mol(smiles)
+            mol_graph = mol_to_bigraph(mol,
+                                       node_featurizer=self.atom_featurizer_enc,
+                                       edge_featurizer=self.bond_featurizer_enc,
+                                       canonical_atom_order=False)
+            mol_graph.apply_edges(fn.copy_u('x', 'src'))
+            mol_graph.edata['x'] = torch.cat(
+                [mol_graph.edata.pop('src'), mol_graph.edata['x']], dim=1)
+
+            if self.cache:
+                self.trees[idx] = mol_tree
+                self.mol_graphs[idx] = mol_graph
+
+        return mol_tree, mol_graph
+
     def __getitem__(self, idx):
-        """Get the datapoint corresponding to the index.
+        """Get a datapoint corresponding to the index.
 
         Parameters
         ----------
         idx : int
-            Index for the datapoint.
+            ID for the datapoint.
 
         Returns
         -------
-        result: dict
-
-            The dictionary contains the following items:
-
-            * result['mol_tree'] : MolTree
-                The junction tree for the original compound.
-            * result['mol_graph'] : DGLGraph
-                The DGLGraph for the original compound.
-            * result['wid'] : list of int
-                The ids corresponding to the clusters (nodes in the junction tree)
-                in the vocabulary.
-            * result['cand_graphs'] : list of DGLGraph, optional
-                DGLGraphs corresponding to the enumerated candidate molecules.
-                This only exists when ``self.training`` is True.
-            * result['tree_mess_src_e'] : 1D LongTensor, optional
-                The edges in the junction tree corresponding to the edges in
-                ``result['cand_graphs']``. This only exists when ``self.training`` is True.
-            * result['tree_mess_tgt_e'] : 1D LongTensor, optional
-                The edges in ``result['cand_graphs']``.
-                This only exists when ``self.training`` is True.
-            * result['tree_mess_tgt_n'] : 1D LongTensor, optional
-                The nodes in ``result['cand_graphs']``.
-                This only exists when ``self.training`` is True.
-            * result['stereo_cand_graphs'], optional
-                DGLGraphs corresponding to enumerated stereoisomers.
-                This only exists when ``self.training`` is True.
-            * result['stereo_cand_label'], optional
-                A 2-tuple of int. The first element is the index of the ground truth
-                stereoisomer in the enumerated stereoisomers. The second element is
-                the number of the enumerated stereoisomers. This only exists when
-                ``self.training`` is True.
+        MolTree
+            MolTree corresponding to the datapoint.
         """
-        smiles = self.data[idx]
-        mol_tree = DGLMolTree(smiles)
-        mol_tree.recover()
-        mol_tree.assemble()
-
-        wid = _set_node_id(mol_tree, self.vocab)
-
-        # prebuild the molecule graph
-        mol_graph = mol2dgl_enc(mol_tree.smiles,
-                                self.atom_featurizer_enc,
-                                self.bond_featurizer_enc)
-
-        result = {
-            'mol_tree': mol_tree,
-            'mol_graph': mol_graph,
-            'wid': wid
-        }
-
-        if not self.training:
-            return result
-
-        # prebuild the candidate graph list
-        cands = []
-        # Enumerate over clusters in the junction tree
-        for node_id, node in mol_tree.nodes_dict.items():
-            # fill in ground truth
-            if node['label'] not in node['cands']:
-                node['cands'].append(node['label'])
-                node['cand_mols'].append(node['label_mol'])
-
-            if node['is_leaf'] or len(node['cands']) == 1:
-                continue
-            cands.extend([(cand, mol_tree, node_id)
-                         for cand in node['cand_mols']])
-        if len(cands) > 0:
-            cand_graphs, tree_mess_src_e, tree_mess_tgt_e, tree_mess_tgt_n = mol2dgl_dec(
-                cands, self.atom_featurizer_dec, self.bond_featurizer_dec)
+        if self.training:
+            return self.getitem_train(idx)
         else:
-            cand_graphs = []
-            tree_mess_src_e = torch.zeros(0, 2).long()
-            tree_mess_tgt_e = torch.zeros(0, 2).long()
-            tree_mess_tgt_n = torch.zeros(0).long()
-
-        # prebuild the stereoisomers
-        cands = mol_tree.stereo_cands
-        if len(cands) > 1:
-            if mol_tree.smiles_3d not in cands:
-                cands.append(mol_tree.smiles_3d)
-
-            stereo_cand_graphs = [mol2dgl_enc(c, self.atom_featurizer_enc,
-                                              self.bond_featurizer_enc)
-                                  for c in cands]
-            stereo_cand_label = [(cands.index(mol_tree.smiles_3d), len(cands))]
-        else:
-            stereo_cand_graphs = []
-            stereo_cand_label = []
+            return self.getitem_non_train(idx)
 
-        result.update({
-            'cand_graphs': cand_graphs,
-            'tree_mess_src_e': tree_mess_src_e,
-            'tree_mess_tgt_e': tree_mess_tgt_e,
-            'tree_mess_tgt_n': tree_mess_tgt_n,
-            'stereo_cand_graphs': stereo_cand_graphs,
-            'stereo_cand_label': stereo_cand_label,
-            })
-
-        return result
-
-def _unpack_field(examples, field):
-    """Get all values of examples under the specified field.
+class JTVAEZINC(JTVAEDataset):
+    """A ZINC subset used in JTVAE
 
     Parameters
     ----------
-    examples : iterable
-        An iterable of objects.
-    field : str
-        The field for value retrieval.
-
-    Returns
-    -------
-    list
-        A list of values.
+    subset : train
+        The subset to use, which can be one of 'train' and 'test'.
+    vocab : JTVAEVocab
+        Vocabulary for JTVAE.
+    cache : bool
+        Whether to cache the trees to speed up data loading or always construct trees on the fly.
     """
-    return [e[field] for e in examples]
+    def __init__(self, subset, vocab, cache=False):
+        dir = get_download_dir()
+        _url = _get_dgl_url('dataset/jtvae.zip')
+        zip_file_path = '{}/jtvae.zip'.format(dir)
+        download(_url, path=zip_file_path, overwrite=False)
+        extract_archive(zip_file_path, '{}/jtvae'.format(dir))
+
+        if subset == 'train':
+            super(JTVAEZINC, self).__init__(data_file='{}/jtvae/train.txt'.format(dir),
+                                            vocab=vocab, cache=cache)
+        elif subset == 'test':
+            super(JTVAEZINC, self).__init__(data_file='{}/jtvae/test.txt'.format(dir),
+                                            vocab=vocab, cache=cache, training=False)
+        else:
+            raise ValueError("Expect subset to be 'train' or 'test', got {}".format(subset))
 
 class JTVAECollator(object):
     """Collate function for JTVAE.
 
-    JTVAE is introduced in `Junction Tree Variational Autoencoder for Molecular Graph Generation
-    <https://arxiv.org/abs/1802.04364>`__.
-
     Parameters
     ----------
     training : bool
         Whether the collate function is for training or not.
     """
-    def __init__(self, training):
+    def __init__(self, training=True):
         self.training = training
 
-    @staticmethod
-    def _batch_and_set(graphs, flatten):
-        """Batch graphs and set an edge feature
+    def collate_train(self, data):
+        batch_trees, batch_mol_graphs, batch_stereo_cand_graphs = map(list, zip(*data))
 
-        Parameters
-        ----------
-        graphs : list of DGLGraph or list of list of DGLGraph
-            Graphs to batch.
-        flatten : bool
-            If True, ``graphs`` is a list of list of DGLGraphs and it will
-            flatten ``graphs`` before calling dgl.batch.
+        batch_tree_graphs = dgl.batch([tree.graph for tree in batch_trees])
+        batch_mol_graphs = dgl.batch(batch_mol_graphs)
 
-        Returns
-        -------
-        DGLGraph
-            Batched DGLGraph.
-        """
-        if flatten:
-            graphs = [g for f in graphs for g in f]
-        graph_batch = batch(graphs)
-        graph_batch.edata['src_x'] = torch.zeros(graph_batch.num_edges(),
-                                                 graph_batch.ndata['x'].shape[1])
-        return graph_batch
+        # Set batch node ID
+        tot = 0
+        for tree in batch_trees:
+            for node_id in tree.nodes_dict:
+                tree.nodes_dict[node_id]['idx'] = tot
+                tot += 1
 
-    def __call__(self, examples):
-        """Batch multiple datapoints
+        batch_stereo_cand_graphs = list(itertools.chain.from_iterable(batch_stereo_cand_graphs))
+        if len(batch_stereo_cand_graphs) == 0:
+            batch_stereo_cand_graphs = dgl.graph(([], []), idtype=batch_mol_graphs.idtype,
+                                                 device=batch_mol_graphs.device)
+        else:
+            batch_stereo_cand_graphs = dgl.batch(batch_stereo_cand_graphs)
 
-        Parameters
-        ----------
-        examples : list of dict
-            Multiple datapoints.
+        stereo_cand_batch_idx = []
+        stereo_cand_labels = []
+        for i, tree in enumerate(batch_trees):
+            cands = tree.stereo_cands
+            if len(cands) == 1:
+                continue
+            if tree.smiles3D not in cands:
+                cands.append(tree.smiles3D)
+            stereo_cand_batch_idx.extend([i] * len(cands))
+            stereo_cand_labels.append((cands.index(tree.smiles3D), len(cands)))
 
-        Returns
-        -------
-        dict
-            Batched datapoint.
-        """
-        # get list of trees
-        mol_trees = _unpack_field(examples, 'mol_tree')
-        wid = _unpack_field(examples, 'wid')
-        for _wid, mol_tree in zip(wid, mol_trees):
-            mol_tree.g.ndata['wid'] = torch.LongTensor(_wid)
+        if len(stereo_cand_batch_idx) == 0:
+            stereo_cand_batch_idx = torch.zeros(0).long()
+        else:
+            stereo_cand_batch_idx = torch.LongTensor(stereo_cand_batch_idx)
 
-        # TODO: either support pickling or get around ctypes pointers using scipy
-        # batch molecule graphs
-        mol_graphs = _unpack_field(examples, 'mol_graph')
-        mol_graph_batch = self._batch_and_set(mol_graphs, False)
+        return batch_trees, batch_tree_graphs, batch_mol_graphs, stereo_cand_batch_idx, \
+               stereo_cand_labels, batch_stereo_cand_graphs
 
-        result = {
-            'mol_trees': mol_trees,
-            'mol_graph_batch': mol_graph_batch
-        }
+    def collate_non_train(self, data):
+        batch_trees, batch_mol_graphs = map(list, zip(*data))
+        if len(batch_mol_graphs) == 1:
+            batch_tree_graphs = batch_trees[0].graph
+            batch_mol_graphs = batch_mol_graphs[0]
+        else:
+            batch_tree_graphs = dgl.batch([tree.graph for tree in batch_trees])
+            batch_mol_graphs = dgl.batch(batch_mol_graphs)
 
-        if not self.training:
-            return result
+        # Set batch node ID
+        tot = 0
+        for tree in batch_trees:
+            for node_id in tree.nodes_dict:
+                tree.nodes_dict[node_id]['idx'] = tot
+                tot += 1
 
-        # batch candidate graphs
-        cand_graphs = _unpack_field(examples, 'cand_graphs')
-        cand_batch_idx = []
-        tree_mess_src_e = _unpack_field(examples, 'tree_mess_src_e')
-        tree_mess_tgt_e = _unpack_field(examples, 'tree_mess_tgt_e')
-        tree_mess_tgt_n = _unpack_field(examples, 'tree_mess_tgt_n')
+        if len(batch_trees) == 1:
+            batch_trees = batch_trees[0]
 
-        n_graph_nodes = 0
-        n_tree_nodes = 0
-        for i in range(len(cand_graphs)):
-            tree_mess_tgt_e[i] += n_graph_nodes
-            tree_mess_src_e[i] += n_tree_nodes
-            tree_mess_tgt_n[i] += n_graph_nodes
-            n_graph_nodes += sum(g.num_nodes() for g in cand_graphs[i])
-            n_tree_nodes += mol_trees[i].g.num_nodes()
-            cand_batch_idx.extend([i] * len(cand_graphs[i]))
-        tree_mess_tgt_e = torch.cat(tree_mess_tgt_e)
-        tree_mess_src_e = torch.cat(tree_mess_src_e)
-        tree_mess_tgt_n = torch.cat(tree_mess_tgt_n)
+        return batch_trees, batch_tree_graphs, batch_mol_graphs
 
-        cand_graph_batch = self._batch_and_set(cand_graphs, True)
+    def __call__(self, data):
+        """Batch multiple datapoints
 
-        # batch stereoisomers
-        stereo_cand_graphs = _unpack_field(examples, 'stereo_cand_graphs')
-        stereo_cand_batch_idx = []
-        for i in range(len(stereo_cand_graphs)):
-            stereo_cand_batch_idx.extend([i] * len(stereo_cand_graphs[i]))
+        Parameters
+        ----------
+        data : list of tuples
+            Multiple datapoints.
 
-        if len(stereo_cand_batch_idx) > 0:
-            stereo_cand_labels = [
-                (label, length)
-                for ex in _unpack_field(examples, 'stereo_cand_label')
-                for label, length in ex
-            ]
-            stereo_cand_labels, stereo_cand_lengths = zip(*stereo_cand_labels)
-            stereo_cand_graph_batch = self._batch_and_set(stereo_cand_graphs, True)
+        Returns
+        -------
+        list of MolTree
+            Junction trees for a batch of datapoints.
+        DGLGraph
+            Batched graph for the junction trees.
+        DGLGraph
+            Batched graph for the molecular graphs.
+        """
+        if self.training:
+            return self.collate_train(data)
         else:
-            stereo_cand_labels = []
-            stereo_cand_lengths = []
-            stereo_cand_graph_batch = None
-            stereo_cand_batch_idx = []
-
-        result.update({
-            'cand_graph_batch': cand_graph_batch,
-            'cand_batch_idx': cand_batch_idx,
-            'tree_mess_tgt_e': tree_mess_tgt_e,
-            'tree_mess_src_e': tree_mess_src_e,
-            'tree_mess_tgt_n': tree_mess_tgt_n,
-            'stereo_cand_graph_batch': stereo_cand_graph_batch,
-            'stereo_cand_batch_idx': stereo_cand_batch_idx,
-            'stereo_cand_labels': stereo_cand_labels,
-            'stereo_cand_lengths': stereo_cand_lengths,
-            })
-
-        return result
+            return self.collate_non_train(data)
diff --git a/python/dgllife/model/model_zoo/__init__.py b/python/dgllife/model/model_zoo/__init__.py
index 3a6ebbc2..13911aa5 100644
--- a/python/dgllife/model/model_zoo/__init__.py
+++ b/python/dgllife/model/model_zoo/__init__.py
@@ -23,7 +23,7 @@
 try:
     # DGMG requires RDKit support
     from .dgmg import *
-    from .jtnn import *
+    from .jtvae import *
 except ImportError:
     pass
 
diff --git a/python/dgllife/model/model_zoo/jtnn/jtmpn.py b/python/dgllife/model/model_zoo/jtnn/jtmpn.py
deleted file mode 100644
index 72396204..00000000
--- a/python/dgllife/model/model_zoo/jtnn/jtmpn.py
+++ /dev/null
@@ -1,172 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612, W1508, I1101, W0221
-# pylint: disable=redefined-outer-name
-
-import os
-import torch
-import torch.nn as nn
-
-import dgl
-import dgl.function as fn
-from dgl import mean_nodes
-
-ELEM_LIST = ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na',
-             'Ca', 'Fe', 'Al', 'I', 'B', 'K', 'Se', 'Zn', 'H', 'Cu', 'Mn', 'unknown']
-
-ATOM_FDIM = len(ELEM_LIST) + 6 + 5 + 1
-BOND_FDIM = 5
-
-PAPER = os.getenv('PAPER', False)
-
-mpn_loopy_bp_msg = fn.copy_src(src='msg', out='msg')
-mpn_loopy_bp_reduce = fn.sum(msg='msg', out='accum_msg')
-
-class LoopyBPUpdate(nn.Module):
-    def __init__(self, hidden_size):
-        super(LoopyBPUpdate, self).__init__()
-        self.hidden_size = hidden_size
-
-        self.W_h = nn.Linear(hidden_size, hidden_size, bias=False)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_h.reset_parameters()
-
-    def forward(self, node):
-        msg_input = node.data['msg_input']
-        msg_delta = self.W_h(node.data['accum_msg'] + node.data['alpha'])
-        msg = torch.relu(msg_input + msg_delta)
-        return {'msg': msg}
-
-if PAPER:
-    mpn_gather_msg = [
-        fn.copy_edge(edge='msg', out='msg'),
-        fn.copy_edge(edge='alpha', out='alpha')
-    ]
-else:
-    mpn_gather_msg = fn.copy_edge(edge='msg', out='msg')
-
-if PAPER:
-    mpn_gather_reduce = [
-        fn.sum(msg='msg', out='m'),
-        fn.sum(msg='alpha', out='accum_alpha'),
-    ]
-else:
-    mpn_gather_reduce = fn.sum(msg='msg', out='m')
-
-class GatherUpdate(nn.Module):
-    def __init__(self, hidden_size):
-        super(GatherUpdate, self).__init__()
-        self.hidden_size = hidden_size
-
-        self.W_o = nn.Linear(ATOM_FDIM + hidden_size, hidden_size)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_o.reset_parameters()
-
-    def forward(self, node):
-        if PAPER:
-            #m = node['m']
-            m = node.data['m'] + node.data['accum_alpha']
-        else:
-            m = node.data['m'] + node.data['alpha']
-        return {
-            'h': torch.relu(self.W_o(torch.cat([node.data['x'], m], 1))),
-        }
-
-class DGLJTMPN(nn.Module):
-    def __init__(self, hidden_size, depth):
-        nn.Module.__init__(self)
-
-        self.depth = depth
-
-        self.W_i = nn.Linear(ATOM_FDIM + BOND_FDIM, hidden_size, bias=False)
-
-        self.loopy_bp_updater = LoopyBPUpdate(hidden_size)
-        self.gather_updater = GatherUpdate(hidden_size)
-        self.hidden_size = hidden_size
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_i.reset_parameters()
-        self.loopy_bp_updater.reset_parameters()
-        self.gather_updater.reset_parameters()
-
-    def forward(self, cand_batch, mol_tree_batch):
-        cand_graphs, tree_mess_src_edges, tree_mess_tgt_edges, tree_mess_tgt_nodes = cand_batch
-
-        cand_line_graph = dgl.line_graph(cand_graphs, backtracking=False, shared=True)
-        cand_line_graph._node_frames = cand_graphs._edge_frames
-
-        cand_graphs = self.run(
-            cand_graphs, cand_line_graph, tree_mess_src_edges, tree_mess_tgt_edges,
-            tree_mess_tgt_nodes, mol_tree_batch)
-
-        g_repr = mean_nodes(cand_graphs, 'h')
-
-        return g_repr
-
-    def run(self, cand_graphs, cand_line_graph, tree_mess_src_edges, tree_mess_tgt_edges,
-            tree_mess_tgt_nodes, mol_tree_batch):
-        device = cand_graphs.device
-        n_nodes = cand_graphs.num_nodes()
-
-        cand_graphs.apply_edges(
-            func=lambda edges: {'src_x': edges.src['x']},
-        )
-
-        bond_features = cand_line_graph.ndata['x']
-        source_features = cand_line_graph.ndata['src_x']
-        features = torch.cat([source_features, bond_features], 1)
-        msg_input = self.W_i(features)
-        cand_line_graph.ndata.update({
-            'msg_input': msg_input,
-            'msg': torch.relu(msg_input),
-            'accum_msg': torch.zeros_like(msg_input),
-        })
-        zero_node_state = bond_features.new(n_nodes, self.hidden_size).zero_()
-        cand_graphs.ndata.update({
-            'm': zero_node_state.clone(),
-            'h': zero_node_state.clone(),
-        })
-
-        cand_graphs.edata['alpha'] = \
-            torch.zeros(cand_graphs.num_edges(), self.hidden_size).to(device)
-        cand_graphs.ndata['alpha'] = zero_node_state
-        if tree_mess_src_edges.shape[0] > 0:
-            if PAPER:
-                src_u, src_v = tree_mess_src_edges.unbind(1)
-                tgt_u, tgt_v = tree_mess_tgt_edges.unbind(1)
-                alpha = mol_tree_batch.edges[src_u, src_v].data['m']
-                cand_graphs.edges[tgt_u, tgt_v].data['alpha'] = alpha
-            else:
-                src_u, src_v = tree_mess_src_edges.unbind(1)
-                alpha = mol_tree_batch.edges[src_u, src_v].data['m']
-                node_idx = (tree_mess_tgt_nodes
-                            .to(device=zero_node_state.device, dtype=torch.int64)[:, None]
-                            .expand_as(alpha))
-                node_alpha = zero_node_state.clone().scatter_add(0, node_idx, alpha)
-                cand_graphs.ndata['alpha'] = node_alpha
-                cand_graphs.apply_edges(
-                    func=lambda edges: {'alpha': edges.src['alpha']},
-                )
-
-        for i in range(self.depth - 1):
-            cand_line_graph.update_all(
-                mpn_loopy_bp_msg,
-                mpn_loopy_bp_reduce,
-                self.loopy_bp_updater,
-            )
-
-        cand_graphs.update_all(
-            mpn_gather_msg,
-            mpn_gather_reduce,
-            self.gather_updater,
-        )
-
-        return cand_graphs
diff --git a/python/dgllife/model/model_zoo/jtnn/jtnn_dec.py b/python/dgllife/model/model_zoo/jtnn/jtnn_dec.py
deleted file mode 100644
index 7efc687b..00000000
--- a/python/dgllife/model/model_zoo/jtnn/jtnn_dec.py
+++ /dev/null
@@ -1,418 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612, W0221, E1102
-
-import dgl
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-import dgl.function as fn
-from dgl import batch, dfs_labeled_edges_generator
-
-from .nnutils import GRUUpdate
-from ....data.jtvae import get_mol, enum_assemble_nx, DGLMolTree
-
-def dfs_order(forest, roots):
-    edges = dfs_labeled_edges_generator(forest, roots, has_reverse_edge=True)
-    for e, l in zip(*edges):
-        # I exploited the fact that the reverse edge ID equal to 1 xor forward
-        # edge ID for molecule trees.  Normally, I should locate reverse edges
-        # using find_edges().
-        yield e ^ l, l
-
-dec_tree_node_msg = fn.copy_edge(edge='m', out='m')
-dec_tree_node_reduce = fn.sum(msg='m', out='h')
-
-def dec_tree_node_update(nodes):
-    return {'new': nodes.data['new'].clone().zero_()}
-
-dec_tree_edge_msg1 = fn.copy_src(src='m', out='m')
-dec_tree_edge_msg2 = fn.copy_src(src='rm', out='rm')
-dec_tree_edge_reduce1 = fn.sum(msg='m', out='s')
-dec_tree_edge_reduce2 = fn.sum(msg='rm', out='accum_rm')
-
-def have_slots(fa_slots, ch_slots):
-    if len(fa_slots) > 2 and len(ch_slots) > 2:
-        return True
-    matches = []
-    for i, s1 in enumerate(fa_slots):
-        a1, c1, h1 = s1
-        for j, s2 in enumerate(ch_slots):
-            a2, c2, h2 = s2
-            if a1 == a2 and c1 == c2 and (a1 != "C" or h1 + h2 >= 4):
-                matches.append((i, j))
-
-    if len(matches) == 0:
-        return False
-
-    fa_match, ch_match = list(zip(*matches))
-    if len(set(fa_match)) == 1 and 1 < len(fa_slots) <= 2:  # never remove atom from ring
-        fa_slots.pop(fa_match[0])
-    if len(set(ch_match)) == 1 and 1 < len(ch_slots) <= 2:  # never remove atom from ring
-        ch_slots.pop(ch_match[0])
-
-    return True
-
-def can_assemble(mol_tree, u, v_node_dict):
-    u_node_dict = mol_tree.nodes_dict[u]
-    u_neighbors = mol_tree.g.successors(u)
-    u_neighbors_node_dict = [
-        mol_tree.nodes_dict[_u]
-        for _u in u_neighbors
-        if _u in mol_tree.nodes_dict
-    ]
-    neis = u_neighbors_node_dict + [v_node_dict]
-    for i, nei in enumerate(neis):
-        nei['nid'] = i
-
-    neighbors = [nei for nei in neis if nei['mol'].GetNumAtoms() > 1]
-    neighbors = sorted(
-        neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
-    singletons = [nei for nei in neis if nei['mol'].GetNumAtoms() == 1]
-    neighbors = singletons + neighbors
-    cands = enum_assemble_nx(u_node_dict, neighbors)
-    return len(cands) > 0
-
-def create_node_dict(smiles, clique=None):
-    if clique is None:
-        clique = []
-    return dict(
-        smiles=smiles,
-        mol=get_mol(smiles),
-        clique=clique,
-    )
-
-class DGLJTNNDecoder(nn.Module):
-    def __init__(self, vocab, hidden_size, latent_size, embedding):
-        nn.Module.__init__(self)
-
-        self.hidden_size = hidden_size
-        self.vocab_size = vocab.size()
-        self.vocab = vocab
-
-        self.embedding = embedding
-        self.dec_tree_edge_update = GRUUpdate(hidden_size)
-
-        self.W = nn.Linear(latent_size + hidden_size, hidden_size)
-        self.U = nn.Linear(latent_size + 2 * hidden_size, hidden_size)
-        self.W_o = nn.Linear(hidden_size, self.vocab_size)
-        self.U_s = nn.Linear(hidden_size, 1)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.embedding.reset_parameters()
-        self.dec_tree_edge_update.reset_parameters()
-        self.W.reset_parameters()
-        self.U.reset_parameters()
-        self.W_o.reset_parameters()
-        self.U_s.reset_parameters()
-
-    def forward(self, mol_trees, tree_vec):
-        '''
-        The training procedure which computes the prediction loss given the
-        ground truth tree
-        '''
-        mol_tree_batch = batch(mol_trees)
-        mol_tree_batch_lg = dgl.line_graph(mol_tree_batch, backtracking=False, shared=True)
-        mol_tree_batch_lg._node_frames = mol_tree_batch._edge_frames
-        n_trees = len(mol_trees)
-
-        return self.run(mol_tree_batch, mol_tree_batch_lg, n_trees, tree_vec)
-
-    def run(self, mol_tree_batch, mol_tree_batch_lg, n_trees, tree_vec):
-        device = tree_vec.device
-
-        node_offset = np.cumsum([0] + mol_tree_batch.batch_num_nodes)
-        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': torch.zeros(n_nodes, self.hidden_size).to(device),
-            # whether it's newly generated node
-            'new': torch.ones(n_nodes).bool().to(device),
-        })
-
-        mol_tree_batch.edata.update({
-            's': torch.zeros(n_edges, self.hidden_size).to(device),
-            'm': torch.zeros(n_edges, self.hidden_size).to(device),
-            'r': torch.zeros(n_edges, self.hidden_size).to(device),
-            'z': torch.zeros(n_edges, self.hidden_size).to(device),
-            'src_x': torch.zeros(n_edges, self.hidden_size).to(device),
-            'dst_x': torch.zeros(n_edges, self.hidden_size).to(device),
-            'rm': torch.zeros(n_edges, self.hidden_size).to(device),
-            'accum_rm': torch.zeros(n_edges, self.hidden_size).to(device),
-        })
-
-        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,
-            dec_tree_node_msg,
-            dec_tree_node_reduce,
-            dec_tree_node_update,
-        )
-        # 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):
-            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(p_target_list.clone().detach())
-
-            root_out_degrees -= (root_out_degrees == 0).long()
-            root_out_degrees -= torch.tensor(np.isin(root_ids,
-                                                     v).astype('int64'))
-
-            mol_tree_batch_lg.pull(
-                eid,
-                dec_tree_edge_msg1,
-                dec_tree_edge_reduce1
-            )
-            mol_tree_batch_lg.pull(
-                eid,
-                dec_tree_edge_msg2,
-                dec_tree_edge_reduce2,
-                self.dec_tree_edge_update
-            )
-
-            is_new = mol_tree_batch.nodes[v].data['new']
-            mol_tree_batch.pull(
-                v,
-                dec_tree_node_msg,
-                dec_tree_node_reduce,
-                dec_tree_node_update,
-            )
-            # 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()).long())
-
-        # Batch compute the stop/label prediction losses
-        p_inputs = torch.cat(p_inputs, 0)
-        p_targets = torch.cat(p_targets, 0).to(device)
-        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(), reduction='sum'
-        ) / n_trees
-        q_loss = F.cross_entropy(q, q_targets, reduction='sum') / 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
-
-    def decode(self, mol_vec, max_decode_len=100):
-        assert mol_vec.shape[0] == 1
-        device = mol_vec.device
-
-        mol_tree = DGLMolTree(None)
-        mol_tree.g = mol_tree.g.to(device)
-
-        init_hidden = torch.zeros(1, self.hidden_size).to(device)
-
-        root_hidden = torch.cat([init_hidden, mol_vec], 1)
-        root_hidden = F.relu(self.W(root_hidden))
-        root_score = self.W_o(root_hidden)
-        _, root_wid = torch.max(root_score, 1)
-        root_wid = root_wid.view(1)
-
-        mol_tree.g.add_nodes(1)   # root
-        mol_tree.g.nodes[0].data['wid'] = root_wid
-        mol_tree.g.nodes[0].data['x'] = self.embedding(root_wid)
-        mol_tree.g.nodes[0].data['h'] = init_hidden
-        mol_tree.g.nodes[0].data['fail'] = torch.tensor([0]).to(device)
-        mol_tree.nodes_dict[0] = root_node_dict = create_node_dict(
-            self.vocab.get_smiles(root_wid))
-
-        stack, trace = [], []
-        stack.append((0, self.vocab.get_slots(root_wid)))
-
-        all_nodes = {0: root_node_dict}
-        first = True
-        new_node_id = 0
-        new_edge_id = 0
-
-        for step in range(max_decode_len):
-            u, u_slots = stack[-1]
-            udata = mol_tree.g.nodes[u].data
-            x = udata['x']
-            h = udata['h']
-
-            # Predict stop
-            p_input = torch.cat([x, h, mol_vec], 1)
-            p_score = torch.sigmoid(self.U_s(torch.relu(self.U(p_input))))
-            backtrack = (p_score.item() < 0.5)
-
-            if not backtrack:
-                # Predict next clique.  Note that the prediction may fail due
-                # to lack of assemblable components
-                mol_tree.g.add_nodes(1)
-                new_node_id += 1
-                v = new_node_id
-                mol_tree.g.add_edges(u, v)
-                uv = new_edge_id
-                new_edge_id += 1
-
-                if first:
-                    mol_tree.g.edata.update({
-                        's': torch.zeros(1, self.hidden_size).to(device),
-                        'm': torch.zeros(1, self.hidden_size).to(device),
-                        'r': torch.zeros(1, self.hidden_size).to(device),
-                        'z': torch.zeros(1, self.hidden_size).to(device),
-                        'src_x': torch.zeros(1, self.hidden_size).to(device),
-                        'dst_x': torch.zeros(1, self.hidden_size).to(device),
-                        'rm': torch.zeros(1, self.hidden_size).to(device),
-                        'accum_rm': torch.zeros(1, self.hidden_size).to(device),
-                    })
-                    first = False
-
-                mol_tree.g.edges[uv].data['src_x'] = mol_tree.g.nodes[u].data['x']
-                # keeping dst_x 0 is fine as h on new edge doesn't depend on that.
-
-                # DGL doesn't dynamically maintain a line graph.
-                mol_tree_lg = dgl.line_graph(mol_tree.g, backtracking=False, shared=True)
-                mol_tree_lg._node_frames = mol_tree.g._edge_frames
-
-                mol_tree_lg.pull(
-                    uv,
-                    dec_tree_edge_msg1,
-                    dec_tree_edge_reduce1
-                )
-                mol_tree_lg.pull(
-                    uv,
-                    dec_tree_edge_msg2,
-                    dec_tree_edge_reduce2,
-                    self.dec_tree_edge_update.update_zm
-                )
-                mol_tree.g.pull(
-                    v,
-                    dec_tree_node_msg,
-                    dec_tree_node_reduce,
-                )
-
-                vdata = mol_tree.g.nodes[v].data
-                h_v = vdata['h']
-                q_input = torch.cat([h_v, mol_vec], 1)
-                q_score = torch.softmax(
-                    self.W_o(torch.relu(self.W(q_input))), -1)
-                _, sort_wid = torch.sort(q_score, 1, descending=True)
-                sort_wid = sort_wid.squeeze()
-
-                next_wid = None
-                for wid in sort_wid.tolist()[:5]:
-                    slots = self.vocab.get_slots(wid)
-                    cand_node_dict = create_node_dict(
-                        self.vocab.get_smiles(wid))
-                    if (have_slots(u_slots, slots) and can_assemble(mol_tree, u, cand_node_dict)):
-                        next_wid = wid
-                        next_slots = slots
-                        next_node_dict = cand_node_dict
-                        break
-
-                if next_wid is None:
-                    # Failed adding an actual children; v is a spurious node
-                    # and we mark it.
-                    vdata['fail'] = torch.tensor([1]).to(device)
-                    backtrack = True
-                else:
-                    next_wid = torch.tensor([next_wid]).to(device)
-                    vdata['wid'] = next_wid
-                    vdata['x'] = self.embedding(next_wid)
-                    mol_tree.nodes_dict[v] = next_node_dict
-                    all_nodes[v] = next_node_dict
-                    stack.append((v, next_slots))
-                    mol_tree.g.add_edges(v, u)
-                    vu = new_edge_id
-                    new_edge_id += 1
-                    mol_tree.g.edges[uv].data['dst_x'] = mol_tree.g.nodes[v].data['x']
-                    mol_tree.g.edges[vu].data['src_x'] = mol_tree.g.nodes[v].data['x']
-                    mol_tree.g.edges[vu].data['dst_x'] = mol_tree.g.nodes[u].data['x']
-
-                    # DGL doesn't dynamically maintain a line graph.
-                    mol_tree_lg = dgl.line_graph(mol_tree.g, backtracking=False, shared=True)
-                    mol_tree_lg._node_frames = mol_tree.g._edge_frames
-                    mol_tree_lg.apply_nodes(
-                        self.dec_tree_edge_update.update_r,
-                        uv
-                    )
-
-            if backtrack:
-                if len(stack) == 1:
-                    break   # At root, terminate
-
-                pu, _ = stack[-2]
-                u_pu = mol_tree.g.edge_ids(u, pu)
-
-                mol_tree_lg.pull(
-                    u_pu,
-                    dec_tree_edge_msg1,
-                    dec_tree_edge_reduce1
-                )
-                mol_tree_lg.pull(
-                    u_pu,
-                    dec_tree_edge_msg2,
-                    dec_tree_edge_reduce2,
-                    self.dec_tree_edge_update
-                )
-                mol_tree.g.pull(
-                    pu,
-                    dec_tree_node_msg,
-                    dec_tree_node_reduce,
-                )
-                stack.pop()
-
-        effective_nodes = mol_tree.g.filter_nodes(
-            lambda nodes: nodes.data['fail'] != 1)
-        effective_nodes, _ = torch.sort(effective_nodes)
-        return mol_tree, all_nodes, effective_nodes
diff --git a/python/dgllife/model/model_zoo/jtnn/jtnn_enc.py b/python/dgllife/model/model_zoo/jtnn/jtnn_enc.py
deleted file mode 100644
index a3bde778..00000000
--- a/python/dgllife/model/model_zoo/jtnn/jtnn_enc.py
+++ /dev/null
@@ -1,140 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612, W0221
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-import dgl
-import dgl.function as fn
-from dgl import bfs_edges_generator
-
-from .nnutils import GRUUpdate
-
-def level_order(forest, roots):
-    device = forest.device
-    edges = list(bfs_edges_generator(forest, roots))
-    edges = [e.to(device) for e in edges]
-    _, leaves = forest.find_edges(edges[-1])
-    edges_back = list(bfs_edges_generator(forest, roots, reverse=True))
-    edges_back = [e.to(device) for e in edges_back]
-    yield from reversed(edges_back)
-    yield from edges
-
-enc_tree_msg1 = fn.copy_src(src='m', out='m')
-enc_tree_msg2 = fn.copy_src(src='rm', out='rm')
-enc_tree_reduce1 = fn.sum(msg='m', out='s')
-enc_tree_reduce2 = fn.sum(msg='rm', out='accum_rm')
-enc_tree_gather_msg = fn.copy_edge(edge='m', out='m')
-enc_tree_gather_reduce = fn.sum(msg='m', out='m')
-
-class EncoderGatherUpdate(nn.Module):
-    def __init__(self, hidden_size):
-        nn.Module.__init__(self)
-        self.hidden_size = hidden_size
-
-        self.W = nn.Linear(2 * hidden_size, hidden_size)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W.reset_parameters()
-
-    def forward(self, nodes):
-        x = nodes.data['x']
-        m = nodes.data['m']
-        return {
-            'h': torch.relu(self.W(torch.cat([x, m], 1))),
-        }
-
-class DGLJTNNEncoder(nn.Module):
-    def __init__(self, vocab, hidden_size, embedding):
-        nn.Module.__init__(self)
-        self.hidden_size = hidden_size
-        self.vocab_size = vocab.size()
-        self.vocab = vocab
-
-        self.embedding = embedding
-        self.enc_tree_update = GRUUpdate(hidden_size)
-        self.enc_tree_gather_update = EncoderGatherUpdate(hidden_size)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.embedding.reset_parameters()
-        self.enc_tree_update.reset_parameters()
-        self.enc_tree_gather_update.reset_parameters()
-
-    def forward(self, mol_tree_batch):
-        # Build line graph to prepare for belief propagation
-        mol_tree_batch_lg = dgl.line_graph(mol_tree_batch, backtracking=False, shared=True)
-        mol_tree_batch_lg._node_frames = mol_tree_batch._edge_frames
-
-        return self.run(mol_tree_batch, mol_tree_batch_lg)
-
-    def run(self, mol_tree_batch, mol_tree_batch_lg):
-        device = mol_tree_batch.device
-
-        # Since tree roots are designated to 0.  In the batched graph we can
-        # simply find the corresponding node ID by looking at node_offset
-        node_offset = np.cumsum([0] + mol_tree_batch.batch_num_nodes().tolist())
-        root_ids = torch.tensor(node_offset[:-1], device=device, dtype=mol_tree_batch.idtype)
-        n_nodes = mol_tree_batch.num_nodes()
-        n_edges = mol_tree_batch.num_edges()
-
-        # Assign structure embeddings to tree nodes
-        mol_tree_batch.ndata.update({
-            'x': self.embedding(mol_tree_batch.ndata['wid']),
-            'h': torch.zeros(n_nodes, self.hidden_size).to(device),
-        })
-
-        # Initialize the intermediate variables according to Eq (4)-(8).
-        # Also initialize the src_x and dst_x fields.
-        # TODO: context?
-        mol_tree_batch.edata.update({
-            's': torch.zeros(n_edges, self.hidden_size).to(device),
-            'm': torch.zeros(n_edges, self.hidden_size).to(device),
-            'r': torch.zeros(n_edges, self.hidden_size).to(device),
-            'z': torch.zeros(n_edges, self.hidden_size).to(device),
-            'src_x': torch.zeros(n_edges, self.hidden_size).to(device),
-            'dst_x': torch.zeros(n_edges, self.hidden_size).to(device),
-            'rm': torch.zeros(n_edges, self.hidden_size).to(device),
-            'accum_rm': torch.zeros(n_edges, self.hidden_size).to(device),
-        })
-
-        # Send the source/destination node features to edges
-        mol_tree_batch.apply_edges(
-            func=lambda edges: {
-                'src_x': edges.src['x'], 'dst_x': edges.dst['x']},
-        )
-
-        # Message passing
-        # I exploited the fact that the reduce function is a sum of incoming
-        # messages, and the uncomputed messages are zero vectors.  Essentially,
-        # we can always compute s_ij as the sum of incoming m_ij, no matter
-        # if m_ij is actually computed or not.
-        for eid in level_order(mol_tree_batch, root_ids):
-            mol_tree_batch_lg.pull(
-                eid,
-                enc_tree_msg1,
-                enc_tree_reduce1
-            )
-            mol_tree_batch_lg.pull(
-                eid,
-                enc_tree_msg2,
-                enc_tree_reduce2,
-                self.enc_tree_update,
-            )
-
-        # Readout
-        mol_tree_batch.update_all(
-            enc_tree_gather_msg,
-            enc_tree_gather_reduce,
-            self.enc_tree_gather_update,
-        )
-
-        root_vecs = mol_tree_batch.nodes[root_ids].data['h']
-
-        return mol_tree_batch, root_vecs
diff --git a/python/dgllife/model/model_zoo/jtnn/jtnn_vae.py b/python/dgllife/model/model_zoo/jtnn/jtnn_vae.py
deleted file mode 100644
index fda2050a..00000000
--- a/python/dgllife/model/model_zoo/jtnn/jtnn_vae.py
+++ /dev/null
@@ -1,339 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612, C0200, W0221, E1102
-
-import copy
-import rdkit.Chem as Chem
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from dgl import batch, unbatch
-from dgl.data.utils import _get_dgl_url, download, get_download_dir, extract_archive
-
-from .jtmpn import DGLJTMPN
-from .jtnn_dec import DGLJTNNDecoder
-from .jtnn_enc import DGLJTNNEncoder
-from .vocab import Vocab
-from .mpn import DGLMPN
-from ....data.jtvae import set_atommap, decode_stereo, copy_edit_mol, attach_mols_nx, \
-    enum_assemble_nx, get_atom_featurizer_enc, get_bond_featurizer_enc, mol2dgl_enc, \
-    get_atom_featurizer_dec, get_bond_featurizer_dec, mol2dgl_dec
-
-class DGLJTNNVAE(nn.Module):
-    """
-    `Junction Tree Variational Autoencoder for Molecular Graph Generation
-    <https://arxiv.org/abs/1802.04364>`__
-
-    Parameters
-    ----------
-    hidden_size : int
-        Size for hidden representations.
-    latent_size : int
-        Size for latent representations of nodes and edges.
-    depth : int
-        The number of times for message passing.
-    vocab_file : str
-        The path to a file of vocabulary, with one SMILES per line. If not
-        specified, it will use the vocabulary extracted from the ZINC dataset.
-    """
-    def __init__(self, hidden_size, latent_size, depth, vocab_file=None):
-        super(DGLJTNNVAE, self).__init__()
-
-        if vocab_file is None:
-            default_dir = get_download_dir()
-            vocab_file = '{}/jtvae/{}.txt'.format(default_dir, 'vocab')
-            zip_file_path = '{}/jtvae.zip'.format(default_dir)
-            download(_get_dgl_url('dataset/jtvae.zip'), path=zip_file_path)
-            extract_archive(zip_file_path, '{}/jtvae'.format(default_dir))
-
-        with open(vocab_file, 'r') as f:
-            self.vocab = Vocab([x.strip("\r\n ") for x in f])
-
-        self.hidden_size = hidden_size
-        self.latent_size = latent_size
-        self.depth = depth
-
-        self.embedding = nn.Embedding(self.vocab.size(), hidden_size)
-        self.mpn = DGLMPN(hidden_size, depth)
-        self.jtnn = DGLJTNNEncoder(self.vocab, hidden_size, self.embedding)
-        self.decoder = DGLJTNNDecoder(self.vocab, hidden_size, latent_size // 2, self.embedding)
-        self.jtmpn = DGLJTMPN(hidden_size, depth)
-
-        self.T_mean = nn.Linear(hidden_size, latent_size // 2)
-        self.T_var = nn.Linear(hidden_size, latent_size // 2)
-        self.G_mean = nn.Linear(hidden_size, latent_size // 2)
-        self.G_var = nn.Linear(hidden_size, latent_size // 2)
-
-        self.atom_featurizer_enc = get_atom_featurizer_enc()
-        self.bond_featurizer_enc = get_bond_featurizer_enc()
-        self.atom_featurizer_dec = get_atom_featurizer_dec()
-        self.bond_featurizer_dec = get_bond_featurizer_dec()
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.embedding.reset_parameters()
-        self.mpn.reset_parameters()
-        self.jtnn.reset_parameters()
-        self.decoder.reset_parameters()
-        self.jtmpn.reset_parameters()
-        self.T_mean.reset_parameters()
-        self.T_var.reset_parameters()
-        self.G_mean.reset_parameters()
-        self.G_var.reset_parameters()
-
-    def encode(self, mol_batch):
-        mol_graphs = mol_batch['mol_graph_batch']
-        mol_vec = self.mpn(mol_graphs)
-
-        mol_trees = [tree.g for tree in mol_batch['mol_trees']]
-        mol_tree_batch = batch(mol_trees)
-        mol_tree_batch, tree_vec = self.jtnn(mol_tree_batch)
-
-        return mol_tree_batch, tree_vec, mol_vec
-
-    def sample(self, tree_vec, mol_vec, e1=None, e2=None):
-        device = tree_vec.device
-
-        tree_mean = self.T_mean(tree_vec)
-        tree_log_var = -torch.abs(self.T_var(tree_vec))
-        mol_mean = self.G_mean(mol_vec)
-        mol_log_var = -torch.abs(self.G_var(mol_vec))
-
-        epsilon = torch.randn(*tree_mean.shape).to(device) if e1 is None else e1
-        tree_vec = tree_mean + torch.exp(tree_log_var / 2) * epsilon
-        epsilon = torch.randn(*mol_mean.shape).to(device) if e2 is None else e2
-        mol_vec = mol_mean + torch.exp(mol_log_var / 2) * epsilon
-
-        z_mean = torch.cat([tree_mean, mol_mean], 1)
-        z_log_var = torch.cat([tree_log_var, mol_log_var], 1)
-
-        return tree_vec, mol_vec, z_mean, z_log_var
-
-    def forward(self, mol_batch, beta=0, e1=None, e2=None):
-        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(
-            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
-
-    def assm(self, mol_batch, mol_tree_batch, mol_vec):
-        device = mol_vec.device
-
-        cands = [mol_batch['cand_graph_batch'],
-                 mol_batch['tree_mess_src_e'],
-                 mol_batch['tree_mess_tgt_e'],
-                 mol_batch['tree_mess_tgt_n']]
-        cand_vec = self.jtmpn(cands, mol_tree_batch)
-        cand_vec = self.G_mean(cand_vec)
-
-        batch_idx = torch.LongTensor(mol_batch['cand_batch_idx']).to(device)
-        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 = torch.LongTensor([label]).to(device)
-                all_loss.append(
-                    F.cross_entropy(cur_score.view(1, -1), label, reduction='sum'))
-
-        all_loss = sum(all_loss) / len(mol_batch['mol_trees'])
-        return all_loss, acc / cnt
-
-    def stereo(self, mol_batch, mol_vec):
-        device = mol_vec.device
-
-        stereo_cands = mol_batch['stereo_cand_graph_batch']
-        batch_idx = mol_batch['stereo_cand_batch_idx']
-        labels = mol_batch['stereo_cand_labels']
-        lengths = mol_batch['stereo_cand_lengths']
-
-        if len(labels) == 0:
-            # Only one stereoisomer exists; do nothing
-            return torch.tensor(0.).to(device), 1.
-
-        batch_idx = torch.LongTensor(batch_idx).to(device)
-        stereo_cands = self.mpn(stereo_cands)
-        stereo_cands = self.G_mean(stereo_cands)
-        stereo_labels = mol_vec[batch_idx]
-        scores = F.cosine_similarity(stereo_cands, stereo_labels)
-
-        st, acc = 0, 0
-        all_loss = []
-        for label, le in zip(labels, lengths):
-            cur_scores = scores[st:st + le]
-            if cur_scores.data[label].item() >= cur_scores.max().item():
-                acc += 1
-            label = torch.LongTensor([label]).to(device)
-            all_loss.append(
-                F.cross_entropy(cur_scores.view(1, -1), label, reduction='sum'))
-            st += le
-
-        all_loss = sum(all_loss) / len(labels)
-        return all_loss, acc / len(labels)
-
-    def decode(self, tree_vec, mol_vec):
-        device = mol_vec.device
-
-        mol_tree, nodes_dict, effective_nodes = self.decoder.decode(tree_vec)
-        effective_nodes_list = effective_nodes.tolist()
-        nodes_dict = [nodes_dict[v] for v in effective_nodes_list]
-
-        for i, (node_id, node) in enumerate(zip(effective_nodes_list, nodes_dict)):
-            node['idx'] = i
-            node['nid'] = i + 1
-            node['is_leaf'] = True
-            if mol_tree.g.in_degrees(node_id) > 1:
-                node['is_leaf'] = False
-                set_atommap(node['mol'], node['nid'])
-
-        mol_tree_sg = mol_tree.g.subgraph(effective_nodes)
-        mol_tree_msg, _ = self.jtnn(mol_tree_sg)
-        mol_tree_msg = unbatch(mol_tree_msg)[0]
-        mol_tree_msg.nodes_dict = nodes_dict
-
-        cur_mol = copy_edit_mol(nodes_dict[0]['mol'])
-        global_amap = [{}] + [{} for _ in nodes_dict]
-        global_amap[1] = {atom.GetIdx(): atom.GetIdx()
-                          for atom in cur_mol.GetAtoms()}
-
-        cur_mol = self.dfs_assemble(
-            mol_tree_msg, mol_vec, cur_mol, global_amap, [], 0, None)
-        if cur_mol is None:
-            return None
-
-        cur_mol = cur_mol.GetMol()
-        set_atommap(cur_mol)
-        cur_mol = Chem.MolFromSmiles(Chem.MolToSmiles(cur_mol))
-        if cur_mol is None:
-            return None
-
-        smiles_2d = Chem.MolToSmiles(cur_mol)
-        stereo_cands = decode_stereo(smiles_2d)
-        if len(stereo_cands) == 1:
-            return stereo_cands[0]
-        stereo_cand_graphs = [mol2dgl_enc(c, self.atom_featurizer_enc, self.bond_featurizer_enc)
-                              for c in stereo_cands]
-        stereo_cand_graphs = batch(stereo_cand_graphs).to(device)
-        stereo_cand_graphs.edata['src_x'] = torch.zeros(
-            stereo_cand_graphs.num_edges(), stereo_cand_graphs.ndata['x'].shape[1]).to(device)
-        stereo_vecs = self.mpn(stereo_cand_graphs)
-        stereo_vecs = self.G_mean(stereo_vecs)
-        scores = F.cosine_similarity(stereo_vecs, mol_vec)
-        _, max_id = scores.max(0)
-        return stereo_cands[max_id.item()]
-
-    def dfs_assemble(self, mol_tree_msg, mol_vec, cur_mol,
-                     global_amap, fa_amap, cur_node_id, fa_node_id):
-        device = mol_vec.device
-        nodes_dict = mol_tree_msg.nodes_dict
-        fa_node = nodes_dict[fa_node_id] if fa_node_id is not None else None
-        cur_node = nodes_dict[cur_node_id]
-
-        fa_nid = fa_node['nid'] if fa_node is not None else -1
-        prev_nodes = [fa_node] if fa_node is not None else []
-
-        children_node_id = [v for v in mol_tree_msg.successors(cur_node_id).tolist()
-                            if nodes_dict[v]['nid'] != fa_nid]
-        children = [nodes_dict[v] for v in children_node_id]
-        neighbors = [nei for nei in children if nei['mol'].GetNumAtoms() > 1]
-        neighbors = sorted(
-            neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
-        singletons = [nei for nei in children if nei['mol'].GetNumAtoms() == 1]
-        neighbors = singletons + neighbors
-
-        cur_amap = [(fa_nid, a2, a1)
-                    for nid, a1, a2 in fa_amap if nid == cur_node['nid']]
-        cands = enum_assemble_nx(cur_node, neighbors, prev_nodes, cur_amap)
-        if len(cands) == 0:
-            return None
-        cand_smiles, cand_mols, cand_amap = list(zip(*cands))
-
-        cands = [(candmol, mol_tree_msg, cur_node_id) for candmol in cand_mols]
-        cand_graphs, tree_mess_src_edges, tree_mess_tgt_edges, tree_mess_tgt_nodes = \
-            mol2dgl_dec(cands, self.atom_featurizer_dec, self.bond_featurizer_dec)
-        cand_graphs = batch(cand_graphs).to(device)
-        tree_mess_src_edges = tree_mess_src_edges.to(device)
-        tree_mess_tgt_edges = tree_mess_tgt_edges.to(device)
-        tree_mess_tgt_nodes = tree_mess_tgt_nodes.to(device)
-        cand_graphs.edata['src_x'] = torch.zeros(cand_graphs.num_edges(),
-                                                 cand_graphs.ndata['x'].shape[1]).to(device)
-
-        cand_vecs = self.jtmpn(
-            (cand_graphs, tree_mess_src_edges,
-             tree_mess_tgt_edges, tree_mess_tgt_nodes),
-            mol_tree_msg,
-        )
-        cand_vecs = self.G_mean(cand_vecs)
-        mol_vec = mol_vec.squeeze()
-        scores = cand_vecs @ mol_vec
-
-        _, cand_idx = torch.sort(scores, descending=True)
-
-        backup_mol = Chem.RWMol(cur_mol)
-        for i in range(len(cand_idx)):
-            cur_mol = Chem.RWMol(backup_mol)
-            pred_amap = cand_amap[cand_idx[i].item()]
-            new_global_amap = copy.deepcopy(global_amap)
-
-            for nei_id, ctr_atom, nei_atom in pred_amap:
-                if nei_id == fa_nid:
-                    continue
-                new_global_amap[nei_id][nei_atom] = new_global_amap[cur_node['nid']][ctr_atom]
-
-            cur_mol = attach_mols_nx(cur_mol, children, [], new_global_amap)
-            new_mol = cur_mol.GetMol()
-            new_mol = Chem.MolFromSmiles(Chem.MolToSmiles(new_mol))
-
-            if new_mol is None:
-                continue
-
-            result = True
-            for nei_node_id, nei_node in zip(children_node_id, children):
-                if nei_node['is_leaf']:
-                    continue
-                cur_mol = self.dfs_assemble(
-                    mol_tree_msg, mol_vec, cur_mol, new_global_amap, pred_amap,
-                    nei_node_id, cur_node_id)
-                if cur_mol is None:
-                    result = False
-                    break
-
-            if result:
-                return cur_mol
-
-        return None
diff --git a/python/dgllife/model/model_zoo/jtnn/mpn.py b/python/dgllife/model/model_zoo/jtnn/mpn.py
deleted file mode 100644
index c5963d53..00000000
--- a/python/dgllife/model/model_zoo/jtnn/mpn.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612, I1101, W0221
-# pylint: disable=redefined-outer-name
-
-import dgl
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-import dgl.function as fn
-from dgl import mean_nodes
-
-ELEM_LIST = ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na',
-             'Ca', 'Fe', 'Al', 'I', 'B', 'K', 'Se', 'Zn', 'H', 'Cu', 'Mn', 'unknown']
-
-ATOM_FDIM = len(ELEM_LIST) + 6 + 5 + 4 + 1
-BOND_FDIM = 5 + 6
-
-mpn_loopy_bp_msg = fn.copy_src(src='msg', out='msg')
-mpn_loopy_bp_reduce = fn.sum(msg='msg', out='accum_msg')
-
-class LoopyBPUpdate(nn.Module):
-    def __init__(self, hidden_size):
-        super(LoopyBPUpdate, self).__init__()
-        self.hidden_size = hidden_size
-
-        self.W_h = nn.Linear(hidden_size, hidden_size, bias=False)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_h.reset_parameters()
-
-    def forward(self, nodes):
-        msg_input = nodes.data['msg_input']
-        msg_delta = self.W_h(nodes.data['accum_msg'])
-        msg = F.relu(msg_input + msg_delta)
-        return {'msg': msg}
-
-mpn_gather_msg = fn.copy_edge(edge='msg', out='msg')
-mpn_gather_reduce = fn.sum(msg='msg', out='m')
-
-class GatherUpdate(nn.Module):
-    def __init__(self, hidden_size):
-        super(GatherUpdate, self).__init__()
-        self.hidden_size = hidden_size
-
-        self.W_o = nn.Linear(ATOM_FDIM + hidden_size, hidden_size)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_o.reset_parameters()
-
-    def forward(self, nodes):
-        m = nodes.data['m']
-        return {
-            'h': F.relu(self.W_o(torch.cat([nodes.data['x'], m], 1))),
-        }
-
-class DGLMPN(nn.Module):
-    def __init__(self, hidden_size, depth):
-        super(DGLMPN, self).__init__()
-
-        self.depth = depth
-
-        self.W_i = nn.Linear(ATOM_FDIM + BOND_FDIM, hidden_size, bias=False)
-
-        self.loopy_bp_updater = LoopyBPUpdate(hidden_size)
-        self.gather_updater = GatherUpdate(hidden_size)
-        self.hidden_size = hidden_size
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_i.reset_parameters()
-        self.loopy_bp_updater.reset_parameters()
-        self.gather_updater.reset_parameters()
-
-    def forward(self, mol_graph):
-        mol_line_graph = dgl.line_graph(mol_graph, backtracking=False, shared=True)
-        mol_line_graph._node_frames = mol_graph._edge_frames
-
-        mol_graph = self.run(mol_graph, mol_line_graph)
-
-        # TODO: replace with unbatch or readout
-        g_repr = mean_nodes(mol_graph, 'h')
-
-        return g_repr
-
-    def run(self, mol_graph, mol_line_graph):
-        n_nodes = mol_graph.num_nodes()
-
-        mol_graph.apply_edges(
-            func=lambda edges: {'src_x': edges.src['x']},
-        )
-
-        e_repr = mol_line_graph.ndata
-        bond_features = e_repr['x']
-        source_features = e_repr['src_x']
-
-        features = torch.cat([source_features, bond_features], 1)
-        msg_input = self.W_i(features)
-        mol_line_graph.ndata.update({
-            'msg_input': msg_input,
-            'msg': F.relu(msg_input),
-            'accum_msg': torch.zeros_like(msg_input),
-        })
-        mol_graph.ndata.update({
-            'm': bond_features.new(n_nodes, self.hidden_size).zero_(),
-            'h': bond_features.new(n_nodes, self.hidden_size).zero_(),
-        })
-
-        for i in range(self.depth - 1):
-            mol_line_graph.update_all(
-                mpn_loopy_bp_msg,
-                mpn_loopy_bp_reduce,
-                self.loopy_bp_updater,
-            )
-
-        mol_graph.edata.update({
-            'msg': mol_line_graph.ndata['msg']
-        })
-        mol_graph.update_all(
-            mpn_gather_msg,
-            mpn_gather_reduce,
-            self.gather_updater,
-        )
-
-        return mol_graph
diff --git a/python/dgllife/model/model_zoo/jtnn/nnutils.py b/python/dgllife/model/model_zoo/jtnn/nnutils.py
deleted file mode 100644
index 620a8438..00000000
--- a/python/dgllife/model/model_zoo/jtnn/nnutils.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612, W0221
-
-import torch
-import torch.nn as nn
-
-class GRUUpdate(nn.Module):
-    def __init__(self, hidden_size):
-        nn.Module.__init__(self)
-        self.hidden_size = hidden_size
-
-        self.W_z = nn.Linear(2 * hidden_size, hidden_size)
-        self.W_r = nn.Linear(hidden_size, hidden_size, bias=False)
-        self.U_r = nn.Linear(hidden_size, hidden_size)
-        self.W_h = nn.Linear(2 * hidden_size, hidden_size)
-
-    def reset_parameters(self):
-        """Reinitialize model parameters."""
-        self.W_z.reset_parameters()
-        self.W_r.reset_parameters()
-        self.U_r.reset_parameters()
-        self.W_h.reset_parameters()
-
-    def update_zm(self, node):
-        src_x = node.data['src_x']
-        s = node.data['s']
-        rm = node.data['accum_rm']
-        z = torch.sigmoid(self.W_z(torch.cat([src_x, s], 1)))
-        m = torch.tanh(self.W_h(torch.cat([src_x, rm], 1)))
-        m = (1 - z) * s + z * m
-        return {'m': m, 'z': z}
-
-    def update_r(self, node, zm=None):
-        dst_x = node.data['dst_x']
-        m = node.data['m'] if zm is None else zm['m']
-        r_1 = self.W_r(dst_x)
-        r_2 = self.U_r(m)
-        r = torch.sigmoid(r_1 + r_2)
-        return {'r': r, 'rm': r * m}
-
-    def forward(self, node):
-        dic = self.update_zm(node)
-        dic.update(self.update_r(node, zm=dic))
-        return dic
diff --git a/python/dgllife/model/model_zoo/jtnn/vocab.py b/python/dgllife/model/model_zoo/jtnn/vocab.py
deleted file mode 100644
index d0a70149..00000000
--- a/python/dgllife/model/model_zoo/jtnn/vocab.py
+++ /dev/null
@@ -1,55 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# pylint: disable=C0111, C0103, E1101, W0611, W0612
-
-import copy
-import rdkit.Chem as Chem
-
-def get_slots(smiles):
-    """
-    Parameters
-    ----------
-    smiles : str
-        A SMILES string.
-
-    Returns
-    -------
-    list of 3-tuples
-        Each tuple consists of the symbol, formal charge and the total
-        number of hydrogen atoms attached of each atom.
-    """
-    mol = Chem.MolFromSmiles(smiles)
-    return [(atom.GetSymbol(), atom.GetFormalCharge(), atom.GetTotalNumHs())
-            for atom in mol.GetAtoms()]
-
-class Vocab(object):
-    """Vocabulary for JTVAE
-
-    Parameters
-    ----------
-    smiles_list : list of str
-        List of SMILES strings for vocabulary.
-    """
-    def __init__(self, smiles_list):
-        self.vocab = smiles_list
-        self.vmap = {x: i for i, x in enumerate(self.vocab)}
-        self.slots = [get_slots(smiles) for smiles in self.vocab]
-
-    def get_index(self, smiles):
-        """Return the index corresponding to the SMILES string"""
-        return self.vmap[smiles]
-
-    def get_smiles(self, idx):
-        """Return the SMILES string corresponding to the index"""
-        return self.vocab[idx]
-
-    def get_slots(self, idx):
-        """Get the slots for the SMILES corresponding to the index"""
-        return copy.deepcopy(self.slots[idx])
-
-    def size(self):
-        """Return the size of the vocabulary"""
-        return len(self.vocab)
diff --git a/python/dgllife/model/model_zoo/jtvae.py b/python/dgllife/model/model_zoo/jtvae.py
new file mode 100644
index 00000000..563e8340
--- /dev/null
+++ b/python/dgllife/model/model_zoo/jtvae.py
@@ -0,0 +1,924 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# JTVAE
+
+from functools import partial
+
+import copy
+import dgl
+import dgl.function as fn
+import rdkit.Chem as Chem
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dgl.traversal import bfs_edges_generator, dfs_labeled_edges_generator
+
+from ...data.jtvae import get_atom_featurizer_enc, get_bond_featurizer_enc
+from ...utils.featurizers import ConcatFeaturizer, atom_type_one_hot, atom_degree_one_hot,\
+    atom_formal_charge_one_hot, atom_is_aromatic, bond_type_one_hot, bond_is_in_ring
+from ...utils.jtvae.chemutils import enum_assemble, set_atommap, copy_edit_mol, attach_mols, \
+    decode_stereo, get_mol
+from ...utils.mol_to_graph import mol_to_bigraph
+
+__all__ = ['JTNNVAE']
+
+MAX_NB = 8
+MAX_DECODE_LEN = 100
+
+class GRUMessage(nn.Module):
+    def __init__(self, hidden_size, msg_field='m'):
+        super(GRUMessage, self).__init__()
+
+        self.U_r = nn.Linear(hidden_size, hidden_size)
+        self.msg_field = msg_field
+
+    def forward(self, edges):
+        r_1 = edges.dst['src_x_r']
+        r_2 = self.U_r(edges.src['h'])
+        r = torch.sigmoid(r_1 + r_2)
+
+        return {self.msg_field: r * edges.src['h']}
+
+class GRUUpdate(nn.Module):
+    def __init__(self, hidden_size):
+        super(GRUUpdate, self).__init__()
+
+        self.W_z = nn.Linear(2 * hidden_size, hidden_size)
+        self.W_h = nn.Linear(2 * hidden_size, hidden_size)
+
+    def forward(self, node):
+        z = torch.sigmoid(self.W_z(
+            torch.cat([node.data['src_x'], node.data['sum_h']], dim=1)))
+        h_input = torch.cat([node.data['src_x'], node.data['sum_gated_h']], dim=1)
+        pre_h = torch.tanh(self.W_h(h_input))
+        new_h = (torch.tensor(1.0).to(z.device) - z) * node.data['sum_h'] + z * pre_h
+        return {'h': new_h}
+
+# pylint: disable=R1710
+def level_order(forest, roots):
+    device = forest.device
+    edges = list(bfs_edges_generator(forest, roots))
+    if len(edges) == 0:
+        return []
+    edges = [e.to(device) for e in edges]
+    edges_back = list(bfs_edges_generator(forest, roots, reverse=True))
+    edges_back = [e.to(device) for e in edges_back]
+    yield from reversed(edges_back)
+    yield from edges
+
+def get_root_ids(graphs):
+    # Get the ID of the root nodes, the first node of all trees
+    batch_num_nodes = graphs.batch_num_nodes().cpu()
+    batch_num_nodes = torch.cat([torch.tensor([0]), batch_num_nodes], dim=0)
+    root_ids = torch.cumsum(batch_num_nodes, dim=0)[:-1]
+
+    return root_ids
+
+class JTNNEncoder(nn.Module):
+    def __init__(self, vocab, hidden_size, embedding=None):
+        super(JTNNEncoder, self).__init__()
+        self.hidden_size = hidden_size
+        self.vocab = vocab
+
+        if embedding is None:
+            self.embedding = nn.Embedding(vocab.size(), hidden_size)
+        else:
+            self.embedding = embedding
+
+        self.W_r = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.gru_message = GRUMessage(hidden_size)
+        self.gru_update = GRUUpdate(hidden_size)
+        self.W = nn.Sequential(
+            nn.Linear(2 * hidden_size, hidden_size),
+            nn.ReLU()
+        )
+
+    def forward(self, tree_graphs):
+        device = tree_graphs.device
+        if 'x' not in tree_graphs.ndata:
+            tree_graphs.ndata['x'] = self.embedding(tree_graphs.ndata['wid'])
+        tree_graphs.apply_edges(fn.copy_u('x', 'src_x'))
+        tree_graphs = tree_graphs.local_var()
+
+        line_tree_graphs = dgl.line_graph(tree_graphs, backtracking=False)
+        line_tree_graphs.ndata.update({
+            'src_x': tree_graphs.edata['src_x'],
+            'src_x_r': self.W_r(tree_graphs.edata['src_x']),
+            # Exploit the fact that the reduce function is a sum of incoming messages,
+            # and uncomputed messages are zero vectors.
+            'h': torch.zeros(line_tree_graphs.num_nodes(), self.hidden_size).to(device),
+            'sum_h': torch.zeros(line_tree_graphs.num_nodes(), self.hidden_size).to(device),
+            'sum_gated_h': torch.zeros(line_tree_graphs.num_nodes(), self.hidden_size).to(device)
+        })
+
+        # Get the ID of the root nodes, the first node of all trees
+        root_ids = get_root_ids(tree_graphs)
+
+        for eid in level_order(tree_graphs, root_ids.to(dtype=tree_graphs.idtype)):
+            line_tree_graphs.pull(v=eid, message_func=fn.copy_u('h', 'h_nei'),
+                                  reduce_func=fn.sum('h_nei', 'sum_h'))
+            line_tree_graphs.pull(v=eid, message_func=self.gru_message,
+                                  reduce_func=fn.sum('m', 'sum_gated_h'))
+            line_tree_graphs.apply_nodes(self.gru_update, v=eid)
+
+        # Readout
+        tree_graphs.ndata['h'] = torch.zeros(tree_graphs.num_nodes(), self.hidden_size).to(device)
+        tree_graphs.edata['h'] = line_tree_graphs.ndata['h']
+        root_ids = root_ids.to(device)
+        tree_graphs.pull(v=root_ids.to(dtype=tree_graphs.idtype),
+                         message_func=fn.copy_e('h', 'm'),
+                         reduce_func=fn.sum('m', 'h'))
+        root_vec = torch.cat([
+            tree_graphs.ndata['x'][root_ids],
+            tree_graphs.ndata['h'][root_ids]
+        ], dim=1)
+        root_vec = self.W(root_vec)
+
+        return tree_graphs.edata['h'], root_vec
+
+def have_slots(fa_slots, ch_slots):
+    if len(fa_slots) > 2 and len(ch_slots) > 2:
+        return True
+    matches = []
+    for i, s1 in enumerate(fa_slots):
+        a1, c1, h1 = s1
+        for j, s2 in enumerate(ch_slots):
+            a2, c2, h2 = s2
+            if a1 == a2 and c1 == c2 and (a1 != "C" or h1 + h2 >= 4):
+                matches.append((i, j))
+
+    if len(matches) == 0:
+        return False
+
+    fa_match, ch_match = zip(*matches)
+    if len(set(fa_match)) == 1 and 1 < len(fa_slots) <= 2: #never remove atom from ring
+        fa_slots.pop(fa_match[0])
+    if len(set(ch_match)) == 1 and 1 < len(ch_slots) <= 2: #never remove atom from ring
+        ch_slots.pop(ch_match[0])
+
+    return True
+
+def can_assemble(node_x, node_y):
+    neis = node_x['neighbors'] + [node_y]
+    for i, nei in enumerate(neis):
+        nei['nid'] = i
+
+    neighbors = [nei for nei in neis if nei['mol'].GetNumAtoms() > 1]
+    neighbors = sorted(neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
+    singletons = [nei for nei in neis if nei['mol'].GetNumAtoms() == 1]
+    neighbors = singletons + neighbors
+    cands = enum_assemble(node_x, neighbors)
+    return len(cands) > 0
+
+def dfs_order(forest, roots):
+    edges = dfs_labeled_edges_generator(forest, roots, has_reverse_edge=True)
+    for e, l in zip(*edges):
+        # Exploit the fact that the reverse edge ID equals to 1 xor forward
+        # edge ID. Normally, this should be done using find_edges().
+        yield e ^ l, l
+
+def mol_tree_node(smiles, wid=None, idx=None, nbrs=None):
+    if nbrs is None:
+        nbrs = []
+    return {'smiles': smiles, 'mol': get_mol(smiles), 'wid': wid, 'idx': idx, 'neighbors': nbrs}
+
+def gru_functional(x, h_nei, wz, wr, ur, wh):
+    hidden_size = x.size()[-1]
+    sum_h = h_nei.sum(dim=1)
+    z_input = torch.cat([x, sum_h], dim=1)
+    z = torch.sigmoid(wz(z_input))
+
+    r_1 = wr(x).view(-1, 1, hidden_size)
+    r_2 = ur(h_nei)
+    r = torch.sigmoid(r_1 + r_2)
+
+    gated_h = r * h_nei
+    sum_gated_h = gated_h.sum(dim=1)
+    h_input = torch.cat([x, sum_gated_h], dim=1)
+    pre_h = torch.tanh(wh(h_input))
+    new_h = (1.0 - z) * sum_h + z * pre_h
+    return new_h
+
+class JTNNDecoder(nn.Module):
+    def __init__(self, vocab, hidden_size, latent_size, embedding=None):
+        super(JTNNDecoder, self).__init__()
+        self.hidden_size = hidden_size
+        self.vocab_size = vocab.size()
+        self.vocab = vocab
+
+        if embedding is None:
+            self.embedding = nn.Embedding(self.vocab_size, hidden_size)
+        else:
+            self.embedding = embedding
+
+        # GRU Weights
+        self.W_r = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.gru_message = GRUMessage(hidden_size)
+        self.gru_update = GRUUpdate(hidden_size)
+
+        # Feature Aggregate Weights
+        self.W = nn.Linear(latent_size + hidden_size, hidden_size)
+        self.U = nn.Linear(latent_size + 2 * hidden_size, hidden_size)
+
+        # Output Weights
+        self.W_o = nn.Linear(hidden_size, self.vocab_size)
+        self.U_s = nn.Linear(hidden_size, 1)
+
+        # Loss Functions
+        self.pred_loss = nn.CrossEntropyLoss(reduction='sum')
+        self.stop_loss = nn.BCEWithLogitsLoss(reduction='sum')
+
+    def forward(self, tree_graphs, tree_vec):
+        device = tree_vec.device
+        batch_size = tree_graphs.batch_size
+
+        root_ids = get_root_ids(tree_graphs)
+
+        if 'x' not in tree_graphs.ndata:
+            tree_graphs.ndata['x'] = self.embedding(tree_graphs.ndata['wid'])
+        if 'src_x' not in tree_graphs.edata:
+            tree_graphs.apply_edges(fn.copy_u('x', 'src_x'))
+        tree_graphs = tree_graphs.local_var()
+        tree_graphs.apply_edges(func=lambda edges: {'dst_wid': edges.dst['wid']})
+
+        line_tree_graphs = dgl.line_graph(tree_graphs, backtracking=False, shared=True)
+        line_num_nodes = line_tree_graphs.num_nodes()
+        line_tree_graphs.ndata.update({
+            'src_x_r': self.W_r(line_tree_graphs.ndata['src_x']),
+            # Exploit the fact that the reduce function is a sum of incoming messages,
+            # and uncomputed messages are zero vectors.
+            'h': torch.zeros(line_num_nodes, self.hidden_size).to(device),
+            'vec': dgl.broadcast_edges(tree_graphs, tree_vec),
+            'sum_h': torch.zeros(line_num_nodes, self.hidden_size).to(device),
+            'sum_gated_h': torch.zeros(line_num_nodes, self.hidden_size).to(device)
+        })
+
+        # input tensors for stop prediction (p) and label prediction (q)
+        pred_hiddens, pred_mol_vecs, pred_targets = [], [], []
+        stop_hiddens, stop_targets = [], []
+
+        # Predict root
+        pred_hiddens.append(torch.zeros(batch_size, self.hidden_size).to(device))
+        pred_targets.append(tree_graphs.ndata['wid'][root_ids.to(device)])
+        pred_mol_vecs.append(tree_vec)
+
+        # Traverse the tree and predict on children
+        for eid, p in dfs_order(tree_graphs, root_ids.to(dtype=tree_graphs.idtype)):
+            eid = eid.to(device)
+            p = p.to(device=device, dtype=tree_graphs.idtype)
+
+            # Message passing excluding the target
+            line_tree_graphs.pull(v=eid, message_func=fn.copy_u('h', 'h_nei'),
+                                  reduce_func=fn.sum('h_nei', 'sum_h'))
+            line_tree_graphs.pull(v=eid, message_func=self.gru_message,
+                                  reduce_func=fn.sum('m', 'sum_gated_h'))
+            line_tree_graphs.apply_nodes(self.gru_update, v=eid)
+
+            # Node aggregation including the target
+            # By construction, the edges of the raw graph follow the order of
+            # (i1, j1), (j1, i1), (i2, j2), (j2, i2), ... The order of the nodes
+            # in the line graph corresponds to the order of the edges in the raw graph.
+            eid = eid.long()
+            reverse_eid = torch.bitwise_xor(eid, torch.tensor(1).to(device))
+            cur_o = line_tree_graphs.ndata['sum_h'][eid] + \
+                    line_tree_graphs.ndata['h'][reverse_eid]
+
+            # Gather targets
+            mask = (p == torch.tensor(0).to(device))
+            pred_list = eid[mask]
+            stop_target = torch.tensor(1).to(device) - p
+
+            # Hidden states for stop prediction
+            stop_hidden = torch.cat([line_tree_graphs.ndata['src_x'][eid],
+                                     cur_o, line_tree_graphs.ndata['vec'][eid]], dim=1)
+            stop_hiddens.append(stop_hidden)
+            stop_targets.extend(stop_target)
+
+            #Hidden states for clique prediction
+            if len(pred_list) > 0:
+                pred_mol_vecs.append(line_tree_graphs.ndata['vec'][pred_list])
+                pred_hiddens.append(line_tree_graphs.ndata['h'][pred_list])
+                pred_targets.append(line_tree_graphs.ndata['dst_wid'][pred_list])
+
+        #Last stop at root
+        root_ids = root_ids.to(device)
+        cur_x = tree_graphs.ndata['x'][root_ids]
+        tree_graphs.edata['h'] = line_tree_graphs.ndata['h']
+        tree_graphs.pull(v=root_ids.to(dtype=tree_graphs.idtype),
+                         message_func=fn.copy_e('h', 'm'), reduce_func=fn.sum('m', 'cur_o'))
+        stop_hidden = torch.cat([cur_x, tree_graphs.ndata['cur_o'][root_ids], tree_vec], dim=1)
+        stop_hiddens.append(stop_hidden)
+        stop_targets.extend(torch.zeros(batch_size).to(device))
+
+        # Predict next clique
+        pred_hiddens = torch.cat(pred_hiddens, dim=0)
+        pred_mol_vecs = torch.cat(pred_mol_vecs, dim=0)
+        pred_vecs = torch.cat([pred_hiddens, pred_mol_vecs], dim=1)
+        pred_vecs = F.relu(self.W(pred_vecs))
+        pred_scores = self.W_o(pred_vecs)
+        pred_targets = torch.cat(pred_targets, dim=0)
+
+        pred_loss = self.pred_loss(pred_scores, pred_targets) / batch_size
+        _, preds = torch.max(pred_scores, dim=1)
+        pred_acc = torch.eq(preds, pred_targets).float()
+        pred_acc = torch.sum(pred_acc) / pred_targets.nelement()
+
+        # Predict stop
+        stop_hiddens = torch.cat(stop_hiddens, dim=0)
+        stop_vecs = F.relu(self.U(stop_hiddens))
+        stop_scores = self.U_s(stop_vecs).squeeze()
+        stop_targets = torch.Tensor(stop_targets).to(device)
+
+        stop_loss = self.stop_loss(stop_scores, stop_targets) / batch_size
+        stops = torch.ge(stop_scores, 0).float()
+        stop_acc = torch.eq(stops, stop_targets).float()
+        stop_acc = torch.sum(stop_acc) / stop_targets.nelement()
+
+        return pred_loss, stop_loss, pred_acc.item(), stop_acc.item()
+
+    def decode(self, mol_vec, prob_decode):
+        device = mol_vec.device
+        stack = []
+        init_hidden = torch.zeros(1, self.hidden_size).to(device)
+        zero_pad = torch.zeros(1, 1, self.hidden_size).to(device)
+
+        # Root Prediction
+        root_hidden = torch.cat([init_hidden, mol_vec], dim=1)
+        root_hidden = F.relu(self.W(root_hidden))
+        root_score = self.W_o(root_hidden)
+        _, root_wid = torch.max(root_score, dim=1)
+        root_wid = root_wid.item()
+
+        root = mol_tree_node(smiles=self.vocab.get_smiles(root_wid), wid=root_wid, idx=0)
+        stack.append((root, self.vocab.get_slots(root['wid'])))
+
+        all_nodes = [root]
+        h = {}
+        for step in range(MAX_DECODE_LEN):
+            node_x, fa_slot = stack[-1]
+            cur_h_nei = [h[(node_y['idx'], node_x['idx'])] for node_y in node_x['neighbors']]
+            if len(cur_h_nei) > 0:
+                cur_h_nei = torch.stack(cur_h_nei, dim=0).view(1, -1, self.hidden_size)
+            else:
+                cur_h_nei = zero_pad
+
+            cur_x = torch.LongTensor([node_x['wid']]).to(device)
+            cur_x = self.embedding(cur_x)
+
+            # Predict stop
+            cur_h = cur_h_nei.sum(dim=1)
+            stop_hidden = torch.cat([cur_x, cur_h, mol_vec], dim=1)
+            stop_hidden = F.relu(self.U(stop_hidden))
+            stop_score = torch.sigmoid(self.U_s(stop_hidden) * 20).squeeze()
+
+            if prob_decode:
+                backtrack = (torch.bernoulli(1.0 - stop_score.data)[0] == 1)
+            else:
+                backtrack = (stop_score.item() < 0.5)
+
+            if not backtrack:  # Forward: Predict next clique
+                new_h = gru_functional(cur_x, cur_h_nei, self.gru_update.W_z, self.W_r,
+                                       self.gru_message.U_r, self.gru_update.W_h)
+                pred_hidden = torch.cat([new_h, mol_vec], dim=1)
+                pred_hidden = F.relu(self.W(pred_hidden))
+                pred_score = torch.softmax(self.W_o(pred_hidden) * 20, dim=1)
+                if prob_decode:
+                    sort_wid = torch.multinomial(pred_score.data.squeeze(), 5)
+                else:
+                    _, sort_wid = torch.sort(pred_score, dim=1, descending=True)
+                    sort_wid = sort_wid.data.squeeze()
+
+                next_wid = None
+                for wid in sort_wid[:5]:
+                    slots = self.vocab.get_slots(wid)
+                    node_y = mol_tree_node(smiles=self.vocab.get_smiles(wid))
+                    if have_slots(fa_slot, slots) and can_assemble(node_x, node_y):
+                        next_wid = wid
+                        next_slots = slots
+                        break
+
+                if next_wid is None:
+                    backtrack = True  # No more children can be added
+                else:
+                    node_y = mol_tree_node(smiles=self.vocab.get_smiles(next_wid),
+                                           wid=next_wid, idx=step + 1, nbrs=[node_x])
+                    h[(node_x['idx'], node_y['idx'])] = new_h[0]
+                    stack.append((node_y, next_slots))
+                    all_nodes.append(node_y)
+
+            if backtrack:  # Backtrack, use if instead of else
+                if len(stack) == 1:
+                    break  # At root, terminate
+
+                node_fa, _ = stack[-2]
+                cur_h_nei = [h[(node_y['idx'], node_x['idx'])] for node_y in node_x['neighbors']
+                             if node_y['idx'] != node_fa['idx']]
+                if len(cur_h_nei) > 0:
+                    cur_h_nei = torch.stack(cur_h_nei, dim=0).view(1, -1, self.hidden_size)
+                else:
+                    cur_h_nei = zero_pad
+
+                new_h = gru_functional(cur_x, cur_h_nei, self.gru_update.W_z, self.W_r,
+                                       self.gru_message.U_r, self.gru_update.W_h)
+                h[(node_x['idx'], node_fa['idx'])] = new_h[0]
+                node_fa['neighbors'].append(node_x)
+                stack.pop()
+
+        return root, all_nodes
+
+class MPN(nn.Module):
+    def __init__(self, hidden_size, depth, in_node_feats=39, in_edge_feats=50):
+        super(MPN, self).__init__()
+
+        self.W_i = nn.Linear(in_edge_feats, hidden_size, bias=False)
+        self.W_h = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.W_o = nn.Sequential(
+            nn.Linear(in_node_feats + hidden_size, hidden_size),
+            nn.ReLU()
+        )
+        self.depth = depth
+
+    def forward(self, mol_graph):
+        mol_graph = mol_graph.local_var()
+        line_mol_graph = dgl.line_graph(mol_graph, backtracking=False)
+
+        line_input = self.W_i(mol_graph.edata['x'])
+        line_mol_graph.ndata['msg_input'] = line_input
+        line_mol_graph.ndata['msg'] = F.relu(line_input)
+
+        # Message passing over the line graph
+        for _ in range(self.depth - 1):
+            line_mol_graph.update_all(message_func=fn.copy_u('msg', 'msg'),
+                                      reduce_func=fn.sum('msg', 'nei_msg'))
+            nei_msg = self.W_h(line_mol_graph.ndata['nei_msg'])
+            line_mol_graph.ndata['msg'] = F.relu(line_input + nei_msg)
+
+        # Message passing over the raw graph
+        mol_graph.edata['msg'] = line_mol_graph.ndata['msg']
+        mol_graph.update_all(message_func=fn.copy_e('msg', 'msg'),
+                             reduce_func=fn.sum('msg', 'nei_msg'))
+
+        raw_input = torch.cat([mol_graph.ndata['x'], mol_graph.ndata['nei_msg']], dim=1)
+        mol_graph.ndata['atom_hiddens'] = self.W_o(raw_input)
+
+        # Readout
+        mol_vecs = dgl.mean_nodes(mol_graph, 'atom_hiddens')
+
+        return mol_vecs
+
+def index_select_ND(source, dim, index):
+    index_size = index.size()
+    suffix_dim = source.size()[1:]
+    final_size = index_size + suffix_dim
+    target = source.index_select(dim, index.view(-1))
+    return target.view(final_size)
+
+class JTMPN(nn.Module):
+
+    def __init__(self, hidden_size, depth, in_node_feats=35, in_edge_feats=40):
+        super(JTMPN, self).__init__()
+        self.hidden_size = hidden_size
+        self.depth = depth
+
+        self.W_i = nn.Linear(in_edge_feats, hidden_size, bias=False)
+        self.W_h = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.W_o = nn.Linear(in_node_feats + hidden_size, hidden_size)
+        self.atom_featurizer = ConcatFeaturizer([
+            partial(atom_type_one_hot,
+                    allowable_set=['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na',
+                                   'Ca', 'Fe', 'Al', 'I', 'B', 'K', 'Se', 'Zn', 'H', 'Cu', 'Mn'],
+                    encode_unknown=True),
+            partial(atom_degree_one_hot, allowable_set=[0, 1, 2, 3, 4], encode_unknown=True),
+            partial(atom_formal_charge_one_hot, allowable_set=[-1, -2, 1, 2],
+                    encode_unknown=True),
+            atom_is_aromatic
+        ])
+        self.bond_featurizer = ConcatFeaturizer([bond_type_one_hot, bond_is_in_ring])
+
+    def forward(self, cand_batch, tree_mess, device='cpu'):
+        fatoms, fbonds = [], []
+        in_bonds, all_bonds = [], []
+        # Ensure index 0 is vec(0)
+        mess_dict, all_mess = {}, [torch.zeros(self.hidden_size).to(device)]
+        total_atoms = 0
+        scope = []
+
+        for e, vec in tree_mess.items():
+            mess_dict[e] = len(all_mess)
+            all_mess.append(vec)
+
+        for mol, all_nodes, _ in cand_batch:
+            n_atoms = mol.GetNumAtoms()
+
+            for atom in mol.GetAtoms():
+                fatoms.append(torch.Tensor(self.atom_featurizer(atom)))
+                in_bonds.append([])
+
+            for bond in mol.GetBonds():
+                a1 = bond.GetBeginAtom()
+                a2 = bond.GetEndAtom()
+                x = a1.GetIdx() + total_atoms
+                y = a2.GetIdx() + total_atoms
+                # Here x_nid,y_nid could be 0
+                x_nid, y_nid = a1.GetAtomMapNum(), a2.GetAtomMapNum()
+                x_bid = all_nodes[x_nid - 1]['idx'] if x_nid > 0 else -1
+                y_bid = all_nodes[y_nid - 1]['idx'] if y_nid > 0 else -1
+
+                bfeature = torch.Tensor(self.bond_featurizer(bond))
+
+                b = len(all_mess) + len(all_bonds)  # bond idx offseted by len(all_mess)
+                all_bonds.append((x, y))
+                fbonds.append(torch.cat([fatoms[x], bfeature], 0))
+                in_bonds[y].append(b)
+
+                b = len(all_mess) + len(all_bonds)
+                all_bonds.append((y, x))
+                fbonds.append(torch.cat([fatoms[y], bfeature], 0))
+                in_bonds[x].append(b)
+
+                if x_bid >= 0 and y_bid >= 0 and x_bid != y_bid:
+                    if (x_bid, y_bid) in mess_dict:
+                        mess_idx = mess_dict[(x_bid, y_bid)]
+                        in_bonds[y].append(mess_idx)
+                    if (y_bid, x_bid) in mess_dict:
+                        mess_idx = mess_dict[(y_bid, x_bid)]
+                        in_bonds[x].append(mess_idx)
+
+            scope.append((total_atoms, n_atoms))
+            total_atoms += n_atoms
+
+        total_bonds = len(all_bonds)
+        total_mess = len(all_mess)
+        fatoms = torch.stack(fatoms, 0).to(device)
+        fbonds = torch.stack(fbonds, 0).to(device)
+        agraph = torch.zeros(total_atoms, MAX_NB).long().to(device)
+        bgraph = torch.zeros(total_bonds, MAX_NB).long().to(device)
+        tree_message = torch.stack(all_mess, dim=0)
+
+        for a in range(total_atoms):
+            for i, b in enumerate(in_bonds[a]):
+                if i == MAX_NB:
+                    break
+                agraph[a, i] = b
+
+        for b1 in range(total_bonds):
+            x, y = all_bonds[b1]
+            for i, b2 in enumerate(in_bonds[x]):  # b2 is offseted by len(all_mess)
+                if i == MAX_NB:
+                    break
+                if b2 < total_mess or all_bonds[b2 - total_mess][0] != y:
+                    bgraph[b1, i] = b2
+
+        binput = self.W_i(fbonds)
+        graph_message = F.relu(binput)
+
+        for i in range(self.depth - 1):
+            message = torch.cat([tree_message, graph_message], dim=0)
+            nei_message = index_select_ND(message, 0, bgraph)
+            nei_message = nei_message.sum(dim=1)
+            nei_message = self.W_h(nei_message)
+            graph_message = F.relu(binput + nei_message)
+
+        message = torch.cat([tree_message, graph_message], dim=0)
+        nei_message = index_select_ND(message, 0, agraph)
+        nei_message = nei_message.sum(dim=1)
+        ainput = torch.cat([fatoms, nei_message], dim=1)
+        atom_hiddens = F.relu(self.W_o(ainput))
+
+        mol_vecs = []
+        for st, le in scope:
+            mol_vec = atom_hiddens.narrow(0, st, le).sum(dim=0) / le
+            mol_vecs.append(mol_vec)
+
+        mol_vecs = torch.stack(mol_vecs, dim=0)
+        return mol_vecs
+
+class JTNNVAE(nn.Module):
+    # TODO
+    def __init__(self, vocab, hidden_size, latent_size, depth, stereo=True):
+        super(JTNNVAE, self).__init__()
+        self.vocab = vocab
+        self.hidden_size = hidden_size
+        self.latent_size = latent_size
+        self.depth = depth
+
+        self.embedding = nn.Embedding(self.vocab.size(), hidden_size)
+        self.jtnn = JTNNEncoder(self.vocab, hidden_size, self.embedding)
+        self.jtmpn = JTMPN(hidden_size, depth)
+        self.mpn = MPN(hidden_size, depth)
+        self.decoder = JTNNDecoder(self.vocab, hidden_size, latent_size // 2, self.embedding)
+
+        self.T_mean = nn.Linear(hidden_size, latent_size // 2)
+        self.T_var = nn.Linear(hidden_size, latent_size // 2)
+        self.G_mean = nn.Linear(hidden_size, latent_size // 2)
+        self.G_var = nn.Linear(hidden_size, latent_size // 2)
+
+        self.assm_loss = nn.CrossEntropyLoss(reduction='sum')
+        self.use_stereo = stereo
+        if stereo:
+            self.stereo_loss = nn.CrossEntropyLoss(reduction='sum')
+
+        self.atom_featurizer = get_atom_featurizer_enc()
+        self.bond_featurizer = get_bond_featurizer_enc()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            if param.dim() == 1:
+                nn.init.constant_(param.data, 0)
+            else:
+                nn.init.xavier_normal_(param.data)
+
+    def encode(self, batch_tree_graphs, batch_mol_graphs):
+        tree_mess, tree_vec = self.jtnn(batch_tree_graphs)
+        mol_vec = self.mpn(batch_mol_graphs)
+        return tree_mess, tree_vec, mol_vec
+
+    def forward(self, batch_trees, batch_tree_graphs, batch_mol_graphs, stereo_cand_batch_idx,
+                stereo_cand_labels, batch_stereo_cand_graphs, beta=0):
+        batch_size = batch_tree_graphs.batch_size
+        device = batch_tree_graphs.device
+        tree_mess, tree_vec, mol_vec = self.encode(batch_tree_graphs, batch_mol_graphs)
+
+        tree_mean = self.T_mean(tree_vec)
+        tree_log_var = -torch.abs(self.T_var(tree_vec))  # Following Mueller et al.
+        mol_mean = self.G_mean(mol_vec)
+        mol_log_var = -torch.abs(self.G_var(mol_vec))    # Following Mueller et al.
+
+        z_mean = torch.cat([tree_mean, mol_mean], dim=1)
+        z_log_var = torch.cat([tree_log_var, mol_log_var], dim=1)
+        kl_loss = -0.5 * torch.sum(torch.tensor(1.0).to(device) + z_log_var - z_mean * z_mean -
+                                   torch.exp(z_log_var)) / batch_size
+
+        epsilon = torch.randn(batch_size, self.latent_size // 2).to(device)
+        tree_vec = tree_mean + torch.exp(tree_log_var / 2) * epsilon
+
+        epsilon = torch.randn(batch_size, self.latent_size // 2).to(device)
+        mol_vec = mol_mean + torch.exp(mol_log_var / 2) * epsilon
+
+        word_loss, topo_loss, word_acc, topo_acc = self.decoder(batch_tree_graphs, tree_vec)
+        assm_loss, assm_acc = self.assm(batch_trees, batch_tree_graphs, mol_vec, tree_mess)
+
+        if self.use_stereo:
+            stereo_loss, stereo_acc = self.stereo(stereo_cand_batch_idx, stereo_cand_labels,
+                                                  batch_stereo_cand_graphs, mol_vec)
+        else:
+            stereo_loss, stereo_acc = torch.tensor(0.).to(device), 0
+
+        loss = word_loss + topo_loss + assm_loss + 2 * stereo_loss + beta * kl_loss
+
+        return loss, kl_loss.item(), word_acc, topo_acc, assm_acc, stereo_acc
+
+    def edata_to_dict(self, g, tree_mess):
+        tree_mess_ = dict()
+        src, dst = g.edges()
+        for i, edge in enumerate(tuple(zip(src.tolist(), dst.tolist()))):
+            tree_mess_[edge] = tree_mess[i]
+
+        return tree_mess_
+
+    def assm(self, batch_trees, tree_graphs, mol_vec, tree_mess):
+        device = tree_graphs.device
+
+        cands = []
+        cand_batch_idx = []
+        for i, tree in enumerate(batch_trees):
+            for _, node in tree.nodes_dict.items():
+                # Leaf node's attachment is determined by neighboring node's attachment
+                if node['is_leaf'] or len(node['cands']) == 1:
+                    continue
+                cands.extend([(cand, tree.nodes_dict, node) for cand in node['cand_mols']])
+                cand_batch_idx.extend([i] * len(node['cands']))
+
+        tree_mess = self.edata_to_dict(tree_graphs, tree_mess)
+        cand_vec = self.jtmpn(cands, tree_mess, device)
+        cand_vec = self.G_mean(cand_vec)
+
+        if len(cand_batch_idx) == 0:
+            cand_batch_idx = torch.zeros(0).long().to(device)
+        else:
+            cand_batch_idx = torch.LongTensor(cand_batch_idx).to(device)
+
+        mol_vec = mol_vec[cand_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 = torch.bmm(mol_vec, cand_vec).squeeze()
+
+        cnt, tot, acc = 0, 0, 0
+        all_loss = []
+        for tree in batch_trees:
+            for _, node in tree.nodes_dict.items():
+                num_cands = len(node['cands'])
+                if node['is_leaf'] or num_cands == 1:
+                    continue
+
+                cnt += 1
+                label = node['cands'].index(node['label'])
+                cur_score = scores.narrow(0, tot, num_cands)
+                tot += num_cands
+
+                if cur_score[label].item() >= cur_score.max().item():
+                    acc += 1
+
+                label = torch.LongTensor([label]).to(device)
+                all_loss.append(self.assm_loss(cur_score.view(1, -1), label))
+
+        if len(all_loss) > 0:
+            all_loss = sum(all_loss) / len(batch_trees)
+        else:
+            all_loss = torch.zeros(1).to(device)
+        return all_loss, acc * 1.0 / cnt
+
+    def stereo(self, batch_idx, batch_labels, batch_stereo_cand_graphs, mol_vec):
+        device = batch_stereo_cand_graphs.device
+
+        if len(batch_labels) == 0:
+            return torch.zeros(1).to(device), 1.0
+
+        stereo_cands = self.mpn(batch_stereo_cand_graphs)
+        stereo_cands = self.G_mean(stereo_cands)
+        stereo_labels = mol_vec[batch_idx]
+        scores = nn.CosineSimilarity()(stereo_cands, stereo_labels)
+
+        st, acc = 0, 0
+        all_loss = []
+        for label, le in batch_labels:
+            cur_scores = scores.narrow(0, st, le)
+            if cur_scores[label].item() >= cur_scores.max().item():
+                acc += 1
+            label = torch.LongTensor([label]).to(device)
+            all_loss.append(self.stereo_loss(cur_scores.view(1, -1), label))
+            st += le
+
+        all_loss = sum(all_loss) / len(batch_labels)
+        return all_loss, acc * 1.0 / len(batch_labels)
+
+    def reconstruct(self, tree_graph, mol_graph, prob_decode=False):
+        device = tree_graph.device
+        _, tree_vec, mol_vec = self.encode(tree_graph, mol_graph)
+
+        tree_mean = self.T_mean(tree_vec)
+        tree_log_var = -torch.abs(self.T_var(tree_vec))  # Following Mueller et al.
+        mol_mean = self.G_mean(mol_vec)
+        mol_log_var = -torch.abs(self.G_var(mol_vec))  # Following Mueller et al.
+
+        epsilon = torch.randn(1, self.latent_size // 2).to(device)
+        tree_vec = tree_mean + torch.exp(tree_log_var / 2) * epsilon
+        epsilon = torch.randn(1, self.latent_size // 2).to(device)
+        mol_vec = mol_mean + torch.exp(mol_log_var / 2) * epsilon
+        return self.decode(tree_vec, mol_vec, prob_decode)
+
+    def sample_prior(self, prob_decode=False):
+        tree_vec = torch.randn(1, self.latent_size // 2)
+        mol_vec = torch.randn(1, self.latent_size // 2)
+        return self.decode(tree_vec, mol_vec, prob_decode)
+
+    def sample_eval(self):
+        tree_vec = torch.randn(1, self.latent_size // 2)
+        mol_vec = torch.randn(1, self.latent_size // 2)
+        all_smiles = []
+        for _ in range(100):
+            s = self.decode(tree_vec, mol_vec, prob_decode=True)
+            all_smiles.append(s)
+        return all_smiles
+
+    def decode(self, tree_vec, mol_vec, prob_decode):
+        device = tree_vec.device
+        pred_root, pred_nodes = self.decoder.decode(tree_vec, prob_decode)
+
+        # Mark nid & is_leaf & atommap
+        for i, node in enumerate(pred_nodes):
+            node['nid'] = i + 1
+            node['is_leaf'] = (len(node['neighbors']) == 1)
+            if len(node['neighbors']) > 1:
+                set_atommap(node['mol'], node['nid'])
+
+        src = []
+        dst = []
+        for node in pred_nodes:
+            cur_id = node['idx']
+            for nbr in node['neighbors']:
+                nbr_id = nbr['idx']
+                src.extend([cur_id])
+                dst.extend([nbr_id])
+        if len(src) == 0:
+            tree_graph = dgl.graph((src, dst), idtype=torch.int32, device=device,
+                                   num_nodes=max([node['idx'] + 1 for node in pred_nodes]))
+        else:
+            tree_graph = dgl.graph((src, dst), idtype=torch.int32, device=device)
+        node_ids = torch.LongTensor([node['idx'] for node in pred_nodes]).to(device)
+        node_wid = torch.LongTensor([node['wid'] for node in pred_nodes]).to(device)
+        tree_graph_x = torch.zeros(tree_graph.num_nodes(), self.hidden_size).to(device)
+        tree_graph_x[node_ids] = self.embedding(node_wid)
+        tree_graph.ndata['x'] = tree_graph_x
+        tree_mess = self.jtnn(tree_graph)[0]
+        tree_mess = self.edata_to_dict(tree_graph, tree_mess)
+
+        cur_mol = copy_edit_mol(pred_root['mol'])
+        global_amap = [{}] + [{} for _ in pred_nodes]
+        global_amap[1] = {atom.GetIdx(): atom.GetIdx() for atom in cur_mol.GetAtoms()}
+
+        cur_mol = self.dfs_assemble(tree_mess, mol_vec, pred_nodes, cur_mol, global_amap, [],
+                                    pred_root, None, prob_decode)
+        if cur_mol is None:
+            return None
+
+        cur_mol = cur_mol.GetMol()
+        set_atommap(cur_mol)
+        cur_mol = Chem.MolFromSmiles(Chem.MolToSmiles(cur_mol))
+        if cur_mol is None:
+            return None
+        if not self.use_stereo:
+            return Chem.MolToSmiles(cur_mol)
+
+        smiles2D = Chem.MolToSmiles(cur_mol)
+        stereo_cands = decode_stereo(smiles2D)
+        if len(stereo_cands) == 1:
+            return stereo_cands[0]
+
+        stereo_cand_graphs = []
+        for cand in stereo_cands:
+            cand = get_mol(cand)
+            cg = mol_to_bigraph(cand, node_featurizer=self.atom_featurizer,
+                                edge_featurizer=self.bond_featurizer,
+                                canonical_atom_order=False)
+            cg.apply_edges(fn.copy_u('x', 'src'))
+            cg.edata['x'] = torch.cat([cg.edata.pop('src'), cg.edata['x']], dim=1)
+            stereo_cand_graphs.append(cg)
+        stereo_cand_graphs = dgl.batch(stereo_cand_graphs).to(device)
+
+        stereo_vecs = self.mpn(stereo_cand_graphs)
+        stereo_vecs = self.G_mean(stereo_vecs)
+        scores = nn.CosineSimilarity()(stereo_vecs, mol_vec)
+        _, max_id = scores.max(dim=0)
+        return stereo_cands[max_id.item()]
+
+    def dfs_assemble(self, tree_mess, mol_vec, all_nodes, cur_mol, global_amap, fa_amap,
+                     cur_node, fa_node, prob_decode):
+        fa_nid = fa_node['nid'] if fa_node is not None else -1
+        prev_nodes = [fa_node] if fa_node is not None else []
+
+        children = [nei for nei in cur_node['neighbors'] if nei['nid'] != fa_nid]
+        neighbors = [nei for nei in children if nei['mol'].GetNumAtoms() > 1]
+        neighbors = sorted(neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
+        singletons = [nei for nei in children if nei['mol'].GetNumAtoms() == 1]
+        neighbors = singletons + neighbors
+
+        cur_amap = [(fa_nid, a2, a1) for nid, a1, a2 in fa_amap if nid == cur_node['nid']]
+        cands = enum_assemble(cur_node, neighbors, prev_nodes, cur_amap)
+        if len(cands) == 0:
+            return None
+        _, cand_mols, cand_amap = zip(*cands)
+
+        cands = [(candmol, all_nodes, cur_node) for candmol in cand_mols]
+
+        cand_vecs = self.jtmpn(cands, tree_mess, mol_vec.device)
+        cand_vecs = self.G_mean(cand_vecs)
+        mol_vec = mol_vec.squeeze()
+        scores = torch.mv(cand_vecs, mol_vec) * 20
+
+        if prob_decode:
+            probs = torch.softmax(scores.view(1, -1)).squeeze() + 1e-5  # prevent prob = 0
+            cand_idx = torch.multinomial(probs, probs.numel())
+        else:
+            _, cand_idx = torch.sort(scores, descending=True)
+
+        backup_mol = Chem.RWMol(cur_mol)
+        for i in range(cand_idx.numel()):
+            cur_mol = Chem.RWMol(backup_mol)
+            pred_amap = cand_amap[cand_idx[i].item()]
+            new_global_amap = copy.deepcopy(global_amap)
+
+            for nei_id, ctr_atom, nei_atom in pred_amap:
+                if nei_id == fa_nid:
+                    continue
+                new_global_amap[nei_id][nei_atom] = new_global_amap[cur_node['nid']][ctr_atom]
+
+            # father is already attached
+            cur_mol = attach_mols(cur_mol, children, [], new_global_amap)
+            new_mol = cur_mol.GetMol()
+            new_mol = Chem.MolFromSmiles(Chem.MolToSmiles(new_mol))
+
+            if new_mol is None:
+                continue
+
+            result = True
+            for nei_node in children:
+                if nei_node['is_leaf']:
+                    continue
+                cur_mol = self.dfs_assemble(tree_mess, mol_vec, all_nodes, cur_mol,
+                                            new_global_amap, pred_amap, nei_node,
+                                            cur_node, prob_decode)
+                if cur_mol is None:
+                    result = False
+                    break
+            if result:
+                return cur_mol
+
+        return None
diff --git a/python/dgllife/model/pretrain/__init__.py b/python/dgllife/model/pretrain/__init__.py
index f78b6247..b546d8bf 100644
--- a/python/dgllife/model/pretrain/__init__.py
+++ b/python/dgllife/model/pretrain/__init__.py
@@ -20,6 +20,7 @@
 
 url = {**moleculenet_url, **generative_url, **property_url, **reaction_url}
 
+# pylint: disable=W0702
 def download_and_load_checkpoint(model_name, model, model_postfix,
                                  local_pretrained_path='pre_trained.pth', log=True):
     """Download pretrained model checkpoint
@@ -48,7 +49,10 @@ def download_and_load_checkpoint(model_name, model, model_postfix,
     local_pretrained_path = '_'.join([model_name, local_pretrained_path])
     download(url_to_pretrained, path=local_pretrained_path, log=log)
     checkpoint = torch.load(local_pretrained_path, map_location='cpu')
-    model.load_state_dict(checkpoint['model_state_dict'])
+    try:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    except:
+        model.load_state_dict(checkpoint)
 
     if log:
         print('Pretrained model loaded')
@@ -77,7 +81,8 @@ def load_pretrained(model_name, log=True):
           with a canonical atom order
         * ``'DGMG_ZINC_random'``: A DGMG model pre-trained on ZINC for molecule generation
           with a random atom order
-        * ``'JTNN_ZINC'``: A JTNN model pre-trained on ZINC for molecule generation
+        * ``'JTVAE_ZINC_no_kl'``: A JTVAE pre-trained on ZINC for molecule generation,
+          without KL regularization
         * ``'wln_center_uspto'``: A WLN model pre-trained on USPTO for reaction prediction
         * ``'wln_rank_uspto'``: A WLN model pre-trained on USPTO for candidate product ranking
         * ``'gin_supervised_contextpred'``: A GIN model pre-trained with supervised learning
diff --git a/python/dgllife/model/pretrain/generative_models.py b/python/dgllife/model/pretrain/generative_models.py
index c7c0a6a9..08aa5c5a 100644
--- a/python/dgllife/model/pretrain/generative_models.py
+++ b/python/dgllife/model/pretrain/generative_models.py
@@ -7,10 +7,6 @@
 #
 # Pre-trained generative models.
 
-import os
-
-from dgl.data.utils import get_download_dir, download, _get_dgl_url, extract_archive
-
 __all__ = ['generative_url',
            'create_generative_model']
 
@@ -19,13 +15,15 @@
     'DGMG_ChEMBL_random': 'pre_trained/dgmg_ChEMBL_random.pth',
     'DGMG_ZINC_canonical': 'pre_trained/dgmg_ZINC_canonical.pth',
     'DGMG_ZINC_random': 'pre_trained/dgmg_ZINC_random.pth',
-    'JTNN_ZINC': 'pre_trained/JTNN_ZINC.pth'
+    # JTVAE pre-trained on ZINC without KL regularization
+    'JTVAE_ZINC_no_kl': 'pre_trained/jtvae_ZINC_no_kl.pth'
 }
 
 try:
     # Things requiring RDKit
     from rdkit import Chem
-    from ...model import DGMG, DGLJTNNVAE
+    from ...model import DGMG, JTNNVAE
+    from ...utils import JTVAEVocab
 except ImportError:
     pass
 
@@ -56,17 +54,12 @@ def create_generative_model(model_name):
                     num_prop_rounds=2,
                     dropout=0.2)
 
-    elif model_name == "JTNN_ZINC":
-        default_dir = get_download_dir()
-        vocab_file = '{}/jtvae/{}.txt'.format(default_dir, 'vocab')
-        if not os.path.exists(vocab_file):
-            zip_file_path = '{}/jtvae.zip'.format(default_dir)
-            download(_get_dgl_url('dataset/jtvae.zip'), path=zip_file_path)
-            extract_archive(zip_file_path, '{}/jtvae'.format(default_dir))
-        return DGLJTNNVAE(vocab_file=vocab_file,
-                          depth=3,
-                          hidden_size=450,
-                          latent_size=56)
+    elif model_name.startswith('JTVAE'):
+        vocab = JTVAEVocab()
+        return JTNNVAE(vocab=vocab,
+                       hidden_size=450,
+                       latent_size=56,
+                       depth=3)
 
     else:
         return None
diff --git a/python/dgllife/utils/__init__.py b/python/dgllife/utils/__init__.py
index 74f44279..1a8167a4 100644
--- a/python/dgllife/utils/__init__.py
+++ b/python/dgllife/utils/__init__.py
@@ -9,5 +9,6 @@
 from .eval import *
 from .featurizers import *
 from .io import *
+from .jtvae import *
 from .mol_to_graph import *
 from .splitters import *
diff --git a/python/dgllife/model/model_zoo/jtnn/__init__.py b/python/dgllife/utils/jtvae/__init__.py
similarity index 57%
rename from python/dgllife/model/model_zoo/jtnn/__init__.py
rename to python/dgllife/utils/jtvae/__init__.py
index c3e9d200..29fe0aad 100644
--- a/python/dgllife/model/model_zoo/jtnn/__init__.py
+++ b/python/dgllife/utils/jtvae/__init__.py
@@ -3,6 +3,8 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
-# JTNN module
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# Utils in JTVAE
 
-from .jtnn_vae import DGLJTNNVAE
+from .vocab import *
diff --git a/python/dgllife/utils/jtvae/chemutils.py b/python/dgllife/utils/jtvae/chemutils.py
new file mode 100644
index 00000000..c8c8b504
--- /dev/null
+++ b/python/dgllife/utils/jtvae/chemutils.py
@@ -0,0 +1,564 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# Chemistry utils in JTVAE
+
+from collections import defaultdict
+from rdkit import Chem
+from rdkit.Chem.EnumerateStereoisomers import EnumerateStereoisomers
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import minimum_spanning_tree
+
+def get_mol(smiles):
+    """Construct RDKit molecule object and kekulize it.
+
+    Parameters
+    ----------
+    smiles : str
+        The SMILES string for a molecule.
+
+    Returns
+    -------
+    rdkit.Chem.rdchem.Mol
+        RDKit molecule object for the input SMILES string.
+    """
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+    Chem.Kekulize(mol)
+    return mol
+
+def decode_stereo(smiles2D):
+    """Enumerate possible stereoisomers for a molecule and get corresponding 3D SMILES strings
+
+    Parameters
+    ----------
+    smiles2D : str
+        2D SMILES string for a molecule.
+
+    Returns
+    -------
+    list of str
+        3D SMILES strings for stereoisomers
+    """
+    mol = Chem.MolFromSmiles(smiles2D)
+    dec_isomers = list(EnumerateStereoisomers(mol))
+
+    dec_isomers = [Chem.MolFromSmiles(Chem.MolToSmiles(mol, isomericSmiles=True))
+                   for mol in dec_isomers]
+    smiles3D = [Chem.MolToSmiles(mol, isomericSmiles=True) for mol in dec_isomers]
+
+    chiralN = [atom.GetIdx() for atom in dec_isomers[0].GetAtoms()
+               if int(atom.GetChiralTag()) > 0 and atom.GetSymbol() == "N"]
+    if len(chiralN) > 0:
+        for mol in dec_isomers:
+            for idx in chiralN:
+                mol.GetAtomWithIdx(idx).SetChiralTag(Chem.rdchem.ChiralType.CHI_UNSPECIFIED)
+            smiles3D.append(Chem.MolToSmiles(mol, isomericSmiles=True))
+
+    return smiles3D
+
+def get_smiles(mol):
+    """Get Kekule SMILES for the input molecule.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        An RDKit molecule object.
+
+    Returns
+    -------
+    str
+        Kekule SMILES string corresponding to the input molecule.
+    """
+    return Chem.MolToSmiles(mol, kekuleSmiles=True)
+
+# pylint: disable=W0703
+def sanitize(mol):
+    """Sanitize and Kekulize the input molecule.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        An RDKit molecule object.
+
+    Returns
+    -------
+    mol : rdkit.Chem.rdchem.Mol
+        Sanitized and Kekulized RDKit molecule object.
+    """
+    try:
+        smiles = get_smiles(mol)
+        mol = get_mol(smiles)
+    except Exception:
+        return None
+    return mol
+
+def copy_atom(atom):
+    """Make a deep copy of the input atom object.
+
+    Parameters
+    ----------
+    atom : rdkit.Chem.rdchem.Atom
+        An RDKit atom object.
+
+    Returns
+    -------
+    rdkit.Chem.rdchem.Atom
+        Deep copy of the input RDKit atom object.
+    """
+    new_atom = Chem.Atom(atom.GetSymbol())
+    new_atom.SetFormalCharge(atom.GetFormalCharge())
+    new_atom.SetAtomMapNum(atom.GetAtomMapNum())
+    return new_atom
+
+def copy_edit_mol(mol):
+    """Make a deep copy of the input molecule object for editing.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        An RDKit molecule object.
+
+    Returns
+    -------
+    mol : rdkit.Chem.rdchem.RWMol
+        Deep copy of the input RDKit molecule object for editing.
+    """
+    new_mol = Chem.RWMol(Chem.MolFromSmiles(''))
+    for atom in mol.GetAtoms():
+        new_atom = copy_atom(atom)
+        new_mol.AddAtom(new_atom)
+    for bond in mol.GetBonds():
+        a1 = bond.GetBeginAtom().GetIdx()
+        a2 = bond.GetEndAtom().GetIdx()
+        bt = bond.GetBondType()
+        new_mol.AddBond(a1, a2, bt)
+    return new_mol
+
+def get_clique_mol(mol, atoms):
+    """Make a deep copy of atom-induced molecule fragment.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        An RDKit molecule object.
+    atoms : list of int
+        IDs for the atoms to induce the molecule fragment.
+
+    Returns
+    -------
+    rdkit.Chem.rdchem.Mol
+        Sanitized atom-induced molecule fragment.
+    """
+    smiles = Chem.MolFragmentToSmiles(mol, atoms, kekuleSmiles=True)
+    new_mol = Chem.MolFromSmiles(smiles, sanitize=False)
+    new_mol = copy_edit_mol(new_mol).GetMol()
+    new_mol = sanitize(new_mol) #We assume this is not None
+    return new_mol
+
+def atom_equal(a1, a2):
+    """Check if two atoms are equal in terms of symbol and formal charge.
+
+    Parameters
+    ----------
+    a1 : rdkit.Chem.rdchem.Atom
+        An RDKit atom object.
+    a2 : rdkit.Chem.rdchem.Atom
+        An RDKit atom object.
+
+    Returns
+    -------
+    bool
+        Whether the two atom objects are equivalent.
+    """
+    return a1.GetSymbol() == a2.GetSymbol() and a1.GetFormalCharge() == a2.GetFormalCharge()
+
+def ring_bond_equal(b1, b2, reverse=False):
+    """Check if two bonds are equal.
+
+    Two bonds are equal if the their beginning and end atoms have the same symbol and
+    formal charge. Bond type not considered because all aromatic (so SINGLE matches DOUBLE).
+
+    Parameters
+    ----------
+    b1 : rdkit.Chem.rdchem.Bond
+        An RDKit bond object.
+    b2 : rdkit.Chem.rdchem.Bond
+        An RDKit bond object.
+    reverse : bool
+        Whether to interchange the role of beginning and end atoms of the second
+        bond in comparison.
+
+    Returns
+    -------
+    bool
+        Whether the two bonds are equal.
+    """
+    b1 = (b1.GetBeginAtom(), b1.GetEndAtom())
+    if reverse:
+        b2 = (b2.GetEndAtom(), b2.GetBeginAtom())
+    else:
+        b2 = (b2.GetBeginAtom(), b2.GetEndAtom())
+    return atom_equal(b1[0], b2[0]) and atom_equal(b1[1], b2[1])
+
+def enum_attach(ctr_mol, nei_node, amap, singletons):
+    """Enumerate possible ways to attach a cluster to the central molecule.
+
+    This version records idx mapping between ctr_mol and nei_mol.
+
+    Parameters
+    ----------
+    ctr_mol : rdkit.Chem.rdchem.Mol
+        The central molecule.
+    nei_node : dict
+        A cluster to attach to the central molecule.
+    amap : list of 3-tuples
+        Each tuple consists of the id of the neighboring cluster,
+        the id of the atom in the central molecule and the id of the same atom in the
+        neighboring cluster.
+    singletons : list of int
+        IDs for the neighboring singletons attached.
+
+    Returns
+    -------
+    list
+        Each element is of the form "amap", corresponding to an attachment configuration.
+    """
+    nei_mol, nei_idx = nei_node['mol'], nei_node['nid']
+    att_confs = []
+    # A black list of atoms in the central molecule connected to singletons
+    black_list = [atom_idx for nei_id, atom_idx, _ in amap if nei_id in singletons]
+    ctr_atoms = [atom for atom in ctr_mol.GetAtoms() if atom.GetIdx() not in black_list]
+    ctr_bonds = [bond for bond in ctr_mol.GetBonds()]
+
+    if nei_mol.GetNumBonds() == 0:  # neighbor singleton
+        nei_atom = nei_mol.GetAtomWithIdx(0)
+        used_list = [atom_idx for _, atom_idx, _ in amap]
+        for atom in ctr_atoms:
+            if atom_equal(atom, nei_atom) and atom.GetIdx() not in used_list:
+                new_amap = amap + [(nei_idx, atom.GetIdx(), 0)]
+                att_confs.append(new_amap)
+
+    elif nei_mol.GetNumBonds() == 1:  # neighbor is a bond
+        bond = nei_mol.GetBondWithIdx(0)
+        bond_val = int(bond.GetBondTypeAsDouble())
+        b1, b2 = bond.GetBeginAtom(), bond.GetEndAtom()
+
+        for atom in ctr_atoms:
+            # Optimize if atom is carbon (other atoms may change valence)
+            if atom.GetAtomicNum() == 6 and atom.GetTotalNumHs() < bond_val:
+                continue
+            if atom_equal(atom, b1):
+                new_amap = amap + [(nei_idx, atom.GetIdx(), b1.GetIdx())]
+                att_confs.append(new_amap)
+            elif atom_equal(atom, b2):
+                new_amap = amap + [(nei_idx, atom.GetIdx(), b2.GetIdx())]
+                att_confs.append(new_amap)
+    else:
+        # intersection is an atom
+        for a1 in ctr_atoms:
+            for a2 in nei_mol.GetAtoms():
+                if atom_equal(a1, a2):
+                    # Optimize if atom is carbon (other atoms may change valence)
+                    if a1.GetAtomicNum() == 6 and a1.GetTotalNumHs() + a2.GetTotalNumHs() < 4:
+                        continue
+                    new_amap = amap + [(nei_idx, a1.GetIdx(), a2.GetIdx())]
+                    att_confs.append(new_amap)
+
+        # intersection is an bond
+        if ctr_mol.GetNumBonds() > 1:
+            for b1 in ctr_bonds:
+                for b2 in nei_mol.GetBonds():
+                    if ring_bond_equal(b1, b2):
+                        new_amap = amap + [(nei_idx, b1.GetBeginAtom().GetIdx(),
+                                            b2.GetBeginAtom().GetIdx()),
+                                           (nei_idx, b1.GetEndAtom().GetIdx(),
+                                            b2.GetEndAtom().GetIdx())]
+                        att_confs.append(new_amap)
+
+                    if ring_bond_equal(b1, b2, reverse=True):
+                        new_amap = amap + [(nei_idx, b1.GetBeginAtom().GetIdx(),
+                                            b2.GetEndAtom().GetIdx()),
+                                           (nei_idx, b1.GetEndAtom().GetIdx(),
+                                            b2.GetBeginAtom().GetIdx())]
+                        att_confs.append(new_amap)
+    return att_confs
+
+def attach_mols(ctr_mol, neighbors, prev_nodes, nei_amap):
+    """Attach clusters to a central molecule.
+
+    Parameters
+    ----------
+    ctr_mol : rdkit.Chem.rdchem.Mol
+        The central molecule.
+    neighbors : list of dict
+        Each element contains the information of a neighboring cluster.
+    prev_nodes : list of dict
+        Each element contains the information of a previous cluster.
+    nei_amap : dict
+        nei_amap[nei_id][nei_atom] maps an atom in a neighboring cluster
+        with id nei_id to an atom in the central molecule.
+
+    Returns
+    -------
+    rdkit.Chem.rdchem.Mol
+        The central molecule with clusters attached.
+    """
+    prev_nids = [node['nid'] for node in prev_nodes]
+    for nei_node in prev_nodes + neighbors:
+        nei_id, nei_mol = nei_node['nid'], nei_node['mol']
+        amap = nei_amap[nei_id]
+        for atom in nei_mol.GetAtoms():
+            if atom.GetIdx() not in amap:
+                new_atom = copy_atom(atom)
+                amap[atom.GetIdx()] = ctr_mol.AddAtom(new_atom)
+
+        if nei_mol.GetNumBonds() == 0:
+            nei_atom = nei_mol.GetAtomWithIdx(0)
+            ctr_atom = ctr_mol.GetAtomWithIdx(amap[0])
+            ctr_atom.SetAtomMapNum(nei_atom.GetAtomMapNum())
+        else:
+            for bond in nei_mol.GetBonds():
+                a1 = amap[bond.GetBeginAtom().GetIdx()]
+                a2 = amap[bond.GetEndAtom().GetIdx()]
+                if ctr_mol.GetBondBetweenAtoms(a1, a2) is None:
+                    ctr_mol.AddBond(a1, a2, bond.GetBondType())
+                elif nei_id in prev_nids: #father node overrides
+                    ctr_mol.RemoveBond(a1, a2)
+                    ctr_mol.AddBond(a1, a2, bond.GetBondType())
+    return ctr_mol
+
+def local_attach(ctr_mol, neighbors, prev_nodes, amap_list):
+    """Pre-process and attach clusters to a central molecule.
+
+    Parameters
+    ----------
+    ctr_mol : rdkit.Chem.rdchem.Mol
+        The central molecule.
+    neighbors : list of dict
+        Each element contains the information of a neighboring cluster.
+    prev_nodes : list of dict
+        Each element contains the information of a neighboring cluster.
+    amap_list : list of 3-tuples
+        Each tuple consists of the id of the neighboring cluster,
+        the id of the atom in the central molecule and the id of the same atom in the
+        neighboring cluster.
+
+    Returns
+    -------
+    rdkit.Chem.rdchem.Mol
+        The central molecule with the clusters attached.
+    """
+    ctr_mol = copy_edit_mol(ctr_mol)
+    nei_amap = {nei['nid']: {} for nei in prev_nodes + neighbors}
+
+    for nei_id, ctr_atom, nei_atom in amap_list:
+        nei_amap[nei_id][nei_atom] = ctr_atom
+
+    ctr_mol = attach_mols(ctr_mol, neighbors, prev_nodes, nei_amap)
+    return ctr_mol.GetMol()
+
+def enum_assemble(node, neighbors, prev_nodes=None, prev_amap=None, max_ncand=2000):
+    """Enumerate all possible ways to assemble the central molecule and the neighboring clusters.
+
+    Parameters
+    ----------
+    node : dict
+        The information of the central molecule.
+    neighbors : list of dict
+        Each element contains the information of a neighboring cluster.
+    prev_nodes : list of dict
+        Each element contains the information of a neighboring cluster.
+    prev_amap : list of 3-tuples
+        Each tuple consists of the id of the neighboring cluster,
+        the id of the atom in the central molecule and the id of the same atom in the
+        neighboring cluster.
+    max_ncand : int
+        Maximum number of candidate assemble ways.
+
+    Returns
+    -------
+    list of 3-tuples
+        Each 3-tuple consists of the SMILES, molecule and amap corresponding to a molecule
+        assembled from the central molecule and the neighboring cluster. The amap is a list
+        of tuples. Each tuple consists of the id of the neighboring cluster,
+        the id of the atom in the central molecule and the id of the same atom in the
+        neighboring cluster.
+    """
+    if prev_nodes is None:
+        prev_nodes = []
+    if prev_amap is None:
+        prev_amap = []
+
+    all_attach_confs = []
+    singletons = [nei_node['nid'] for nei_node in neighbors + prev_nodes
+                  if nei_node['mol'].GetNumAtoms() == 1]
+
+    def search(cur_amap, depth):
+        if len(all_attach_confs) > max_ncand:
+            return
+        if depth == len(neighbors):
+            all_attach_confs.append(cur_amap)
+            return
+
+        nei_node = neighbors[depth]
+        cand_amap = enum_attach(node['mol'], nei_node, cur_amap, singletons)
+        cand_smiles = set()
+        candidates = []
+        for amap in cand_amap:
+            cand_mol = local_attach(node['mol'], neighbors[:depth+1], prev_nodes, amap)
+            cand_mol = sanitize(cand_mol)
+            if cand_mol is None:
+                continue
+            smiles = get_smiles(cand_mol)
+            if smiles in cand_smiles:
+                continue
+            cand_smiles.add(smiles)
+            candidates.append(amap)
+
+        if len(candidates) == 0:
+            return
+
+        for new_amap in candidates:
+            search(new_amap, depth + 1)
+
+    search(prev_amap, 0)
+    cand_smiles = set()
+    candidates = []
+    for amap in all_attach_confs:
+        cand_mol = local_attach(node['mol'], neighbors, prev_nodes, amap)
+        cand_mol = Chem.MolFromSmiles(Chem.MolToSmiles(cand_mol))
+        smiles = Chem.MolToSmiles(cand_mol)
+        if smiles in cand_smiles:
+            continue
+        cand_smiles.add(smiles)
+        Chem.Kekulize(cand_mol)
+        candidates.append((smiles, cand_mol, amap))
+
+    return candidates
+
+def set_atommap(mol, num=0):
+    """Set the atom map number for all atoms in the molecule.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        A molecule.
+    num : int
+        The atom map number to set for all atoms. If 0, it will clear the atom map.
+    """
+    for atom in mol.GetAtoms():
+        atom.SetAtomMapNum(num)
+
+# pylint: disable=C0200
+def tree_decomp(mol, mst_max_weight=100):
+    """Tree decomposition of a molecule for junction tree construction.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        A molecule.
+    mst_max_weight : int
+        Max weight considered in generating a minimum spanning tree.
+
+    Returns
+    -------
+    list
+        Clusters. Each element is a list of int,
+        representing the atoms that constitute the cluster.
+    list
+        Edges between the clusters. Each element is a 2-tuple of cluster IDs.
+    """
+    n_atoms = mol.GetNumAtoms()
+    if n_atoms == 1:
+        return [[0]], []
+
+    cliques = []
+
+    # Find all edges not belonging to any cycles
+    for bond in mol.GetBonds():
+        a1 = bond.GetBeginAtom().GetIdx()
+        a2 = bond.GetEndAtom().GetIdx()
+        if not bond.IsInRing():
+            cliques.append([a1, a2])
+
+    # Find all simple cycles, each represented by a list of IDs of the atoms in the ring
+    ssr = [list(x) for x in Chem.GetSymmSSSR(mol)]
+    cliques.extend(ssr)
+
+    # Record the non-ring bonds/simple cycles that each atom belongs to
+    nei_list = [[] for _ in range(n_atoms)]
+    for i in range(len(cliques)):
+        for atom in cliques[i]:
+            nei_list[atom].append(i)
+
+    # Merge simple rings that have more than two overlapping atoms
+    for i in range(len(cliques)):
+        if len(cliques[i]) <= 2:
+            continue
+        for atom in cliques[i]:
+            for j in nei_list[atom]:
+                if i >= j or len(cliques[j]) <= 2:
+                    continue
+                inter = set(cliques[i]) & set(cliques[j])
+                if len(inter) > 2:
+                    cliques[i].extend(cliques[j])
+                    cliques[i] = list(set(cliques[i]))
+                    cliques[j] = []
+
+    # Remove merged simple cycles
+    cliques = [c for c in cliques if len(c) > 0]
+    # Record the non-ring bonds/simple cycles that each atom belongs to
+    nei_list = [[] for _ in range(n_atoms)]
+    for i in range(len(cliques)):
+        for atom in cliques[i]:
+            nei_list[atom].append(i)
+
+    # Build edges and add singleton cliques
+    edges = defaultdict(int)
+    for atom in range(n_atoms):
+        if len(nei_list[atom]) <= 1:
+            continue
+        cnei = nei_list[atom]
+        bonds = [c for c in cnei if len(cliques[c]) == 2]
+        rings = [c for c in cnei if len(cliques[c]) > 4]
+        # In general, if len(cnei) >= 3, a singleton should be added,
+        # but 1 bond + 2 ring is currently not dealt with.
+        if len(bonds) > 2 or (len(bonds) == 2 and len(cnei) > 2):
+            # Add singleton clique
+            cliques.append([atom])
+            c2 = len(cliques) - 1
+            for c1 in cnei:
+                edges[(c1, c2)] = 1
+        elif len(rings) > 2:
+            # Multiple (n>2) complex rings
+            cliques.append([atom])
+            c2 = len(cliques) - 1
+            for c1 in cnei:
+                edges[(c1, c2)] = mst_max_weight - 1
+        else:
+            for i in range(len(cnei)):
+                for j in range(i + 1, len(cnei)):
+                    c1, c2 = cnei[i], cnei[j]
+                    inter = set(cliques[c1]) & set(cliques[c2])
+                    if edges[(c1, c2)] < len(inter):
+                        # cnei[i] < cnei[j] by construction
+                        edges[(c1, c2)] = len(inter)
+
+    edges = [u + (mst_max_weight - v,) for u, v in edges.items()]
+    if len(edges) == 0:
+        return cliques, edges
+
+    # Compute Maximum Spanning Tree
+    row, col, data = zip(*edges)
+    n_clique = len(cliques)
+    clique_graph = csr_matrix((data, (row, col)), shape=(n_clique, n_clique))
+    junc_tree = minimum_spanning_tree(clique_graph)
+    row, col = junc_tree.nonzero()
+    edges = [(row[i], col[i]) for i in range(len(row))]
+    return (cliques, edges)
diff --git a/python/dgllife/utils/jtvae/mol_tree.py b/python/dgllife/utils/jtvae/mol_tree.py
new file mode 100644
index 00000000..3b37b9a9
--- /dev/null
+++ b/python/dgllife/utils/jtvae/mol_tree.py
@@ -0,0 +1,186 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# MolTree in JTVAE
+
+import dgl
+import numpy as np
+import torch
+
+from rdkit import Chem
+
+from .chemutils import get_mol, decode_stereo, get_clique_mol, get_smiles, enum_assemble, \
+    set_atommap, tree_decomp
+
+class MolTree(object):
+    """Junction tree.
+
+    Parameters
+    ----------
+    smiles : str
+        A SMILES string.
+    """
+    def __init__(self, smiles):
+        self.smiles = smiles
+        self.mol = get_mol(smiles)
+        self.nodes_dict = {}
+
+        # Stereo Generation
+        mol = Chem.MolFromSmiles(smiles)
+
+        if mol is None:
+            self.smiles3D = None
+            self.smiles2D = None
+            self.stereo_cands = []
+            self.graph = dgl.graph(([], []), idtype=torch.int32)
+            return
+
+        self.smiles3D = Chem.MolToSmiles(mol, isomericSmiles=True)
+        self.smiles2D = Chem.MolToSmiles(mol)
+        self.stereo_cands = decode_stereo(self.smiles2D)
+
+        # Junction tree construction
+        cliques, edges = tree_decomp(self.mol)
+        root = 0
+        for i, c in enumerate(cliques):
+            cmol = get_clique_mol(self.mol, c)
+            csmiles = get_smiles(cmol)
+            self.nodes_dict[i] = {'smiles': csmiles, 'mol': get_mol(csmiles), 'clique': c}
+
+            if min(c) == 0:
+                root = i
+
+        # Make the clique with atom ID 0 the root
+        root_changed = False
+        if root > 0:
+            root_changed = True
+            for attr in self.nodes_dict[0]:
+                self.nodes_dict[0][attr], self.nodes_dict[root][attr] = \
+                    self.nodes_dict[root][attr], self.nodes_dict[0][attr]
+
+        # Construct DGLGraph for the junction tree
+        src = np.zeros((len(edges) * 2,), dtype='int')
+        dst = np.zeros((len(edges) * 2,), dtype='int')
+
+        def _switch_id(nid, root, root_changed):
+            if not root_changed:
+                return nid
+            if nid == root:
+                return 0
+            elif nid == 0:
+                return root
+            else:
+                return nid
+
+        for i, (_x, _y) in enumerate(edges):
+            x = _switch_id(_x, root, root_changed)
+            y = _switch_id(_y, root, root_changed)
+
+            src[2 * i] = x
+            dst[2 * i] = y
+            src[2 * i + 1] = y
+            dst[2 * i + 1] = x
+
+        self.graph = dgl.graph((src, dst), num_nodes=len(cliques), idtype=torch.int32)
+
+        for i in self.nodes_dict:
+            self.nodes_dict[i]['nid'] = i + 1
+            if self.graph.out_degrees(i) > 1:
+                # Leaf node mol is not marked
+                set_atommap(self.nodes_dict[i]['mol'], self.nodes_dict[i]['nid'])
+            self.nodes_dict[i]['is_leaf'] = (self.graph.out_degrees(i) == 1)
+
+    def size(self):
+        """Get the number of nodes in the junction tree.
+
+        Returns
+        -------
+        int
+            Number of nodes in the junction tree.
+        """
+        return self.graph.num_nodes()
+
+    def _recover_node(self, i, original_mol):
+        """Get the SMILES string corresponding to the i-th cluster in the
+        original molecule.
+
+        Parameters
+        ----------
+        i : int
+            The id of a cluster.
+        original_mol : rdkit.Chem.rdchem.Mol
+            The original molecule.
+
+        Returns
+        -------
+        str
+            A SMILES string.
+        """
+        node = self.nodes_dict[i]
+
+        clique = []
+        clique.extend(node['clique'])
+        if not node['is_leaf']:
+            for cidx in node['clique']:
+                original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(node['nid'])
+
+        for j in self.graph.successors(i).numpy():
+            nei_node = self.nodes_dict[j]
+            clique.extend(nei_node['clique'])
+            if nei_node['is_leaf']:  # Leaf node, no need to mark
+                continue
+            for cidx in nei_node['clique']:
+                # allow singleton node override the atom mapping
+                if cidx not in node['clique'] or len(nei_node['clique']) == 1:
+                    atom = original_mol.GetAtomWithIdx(cidx)
+                    atom.SetAtomMapNum(nei_node['nid'])
+
+        clique = list(set(clique))
+        label_mol = get_clique_mol(original_mol, clique)
+        node['label'] = Chem.MolToSmiles(Chem.MolFromSmiles(get_smiles(label_mol)))
+        node['label_mol'] = get_mol(node['label'])
+
+        for cidx in clique:
+            original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(0)
+
+        return node['label']
+
+    def recover(self):
+        """Get the SMILES string corresponding to all clusters in the original molecule."""
+        for i in self.nodes_dict:
+            self._recover_node(i, self.mol)
+
+    def _assemble_node(self, i):
+        """Assemble a cluster with its successors.
+
+        Parameters
+        ----------
+        i : int
+            The id of a cluster.
+        """
+        neighbors = [self.nodes_dict[j] for j in self.graph.successors(i).numpy()
+                     if self.nodes_dict[j]['mol'].GetNumAtoms() > 1]
+        neighbors = sorted(neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
+        singletons = [self.nodes_dict[j] for j in self.graph.successors(i).numpy()
+                      if self.nodes_dict[j]['mol'].GetNumAtoms() == 1]
+        # All successors sorted based on their size
+        neighbors = singletons + neighbors
+
+        cands = enum_assemble(self.nodes_dict[i], neighbors)
+
+        if len(cands) > 0:
+            cands, cand_mols, _ = list(zip(*cands))
+            self.nodes_dict[i]['cands'] = list(cands)
+            self.nodes_dict[i]['cand_mols'] = list(cand_mols)
+        else:
+            self.nodes_dict[i]['cands'] = []
+            self.nodes_dict[i]['cand_mols'] = []
+
+    def assemble(self):
+        """Assemble each cluster with its successors."""
+        for i in self.nodes_dict:
+            self._assemble_node(i)
diff --git a/python/dgllife/utils/jtvae/vocab.py b/python/dgllife/utils/jtvae/vocab.py
new file mode 100644
index 00000000..d26639a3
--- /dev/null
+++ b/python/dgllife/utils/jtvae/vocab.py
@@ -0,0 +1,112 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# Vocab in JTVAE
+
+from copy import  deepcopy
+from rdkit import Chem
+
+from .mol_tree import MolTree
+
+__all__ = ['JTVAEVocab']
+
+def get_slots(smiles):
+    mol = Chem.MolFromSmiles(smiles)
+    return [(atom.GetSymbol(), atom.GetFormalCharge(),
+             atom.GetTotalNumHs()) for atom in mol.GetAtoms()]
+
+class JTVAEVocab(object):
+    """Vocabulary for JTVAE.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to the vocabulary file, which is a txt file of SMILES strings,
+        one SMILES string a line.
+    """
+    def __init__(self, file_path=None):
+        if file_path is None:
+            from dgl.data.utils import get_download_dir, download, _get_dgl_url, extract_archive
+
+            default_dir = get_download_dir()
+            vocab_file = '{}/jtvae/vocab.txt'.format(default_dir)
+            zip_file_path = '{}/jtvae.zip'.format(default_dir)
+            download(_get_dgl_url('dataset/jtvae.zip'), path=zip_file_path, overwrite=False)
+            extract_archive(zip_file_path, '{}/jtvae'.format(default_dir))
+
+            with open(vocab_file, 'r') as f:
+                self.vocab = [x.strip("\r\n ") for x in f]
+        else:
+            # Prepare a vocabulary from scratch
+            vocab = set()
+            with open(file_path, 'r') as f:
+                for line in f:
+                    smiles = line.split()[0]
+                    mol = MolTree(smiles)
+                    for i in mol.nodes_dict:
+                        vocab.add(mol.nodes_dict[i]['smiles'])
+            self.vocab = list(vocab)
+
+        self.vmap = {x: i for i, x in enumerate(self.vocab)}
+        self.slots = [get_slots(smiles) for smiles in self.vocab]
+
+    def get_index(self, smiles):
+        """Get the index for a vocabulary token.
+
+        Parameters
+        ----------
+        smiles : str
+            The SMILES string for a vocabulary token.
+
+        Returns
+        -------
+        int
+            The ID for the token.
+        """
+        return self.vmap[smiles]
+
+    def get_smiles(self, idx):
+        """Get the SMILES string corresponding to the index.
+
+        Parameters
+        ----------
+        idx : int
+            The ID for a vocabulary token.
+
+        Returns
+        -------
+        str
+            The vocabulary token corresponding to the index.
+        """
+        return self.vocab[idx]
+
+    def get_slots(self, idx):
+        """Get 3-tuples of atom symbol, formal charge and total number of hydrogen atoms for
+        all atoms in the vocabulary token corresponding to the index.
+
+        Parameters
+        ----------
+        idx : int
+            The ID for a vocabulary token.
+
+        Returns
+        -------
+        list of 3-tuples
+            Each tuple consists of symbol, formal charge and total number of hydrogen atoms for
+            an atom in the vocabulary token.
+        """
+        return deepcopy(self.slots[idx])
+
+    def size(self):
+        """Get the size of the vocabulary.
+
+        Returns
+        -------
+        int
+            The vocabulary size.
+        """
+        return len(self.vocab)


Test patch:
diff --git a/tests/data/test_datasets.py b/tests/data/test_datasets.py
index f54764c5..9ae24245 100644
--- a/tests/data/test_datasets.py
+++ b/tests/data/test_datasets.py
@@ -7,11 +7,9 @@
 import shutil
 import torch
 
-from dgl.data.utils import get_download_dir, download, _get_dgl_url, extract_archive
 from dgllife.data import *
 from dgllife.data.uspto import get_bond_changes, process_file
-from dgllife.model.model_zoo.jtnn.vocab import Vocab
-from torch.utils.data import DataLoader
+from dgllife.utils import JTVAEVocab
 
 def remove_file(fname):
     if os.path.isfile(fname):
@@ -192,70 +190,17 @@ def test_wln_reaction():
     remove_file('_invalid_reactions.proc')
 
 def test_jtvae():
-    # Test DGLMolTree
-    smiles = 'CC1([C@@H](N2[C@H](S1)[C@@H](C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C'
-    tree = DGLMolTree(smiles)
-    assert tree.treesize() == 17
-    tree.assemble()
-    assert tree._recover_node(0, tree.mol) == 'C[CH3:15]'
-    tree.recover()
-
-    # Test JTVAEDataset
-    smiles = ['CCCCCCC1=NN2C(=N)/C(=C\c3cc(C)n(-c4ccc(C)cc4C)c3C)C(=O)N=C2S1',
-              'COCC[C@@H](C)C(=O)N(C)Cc1ccc(O)cc1']
-    with open('data.txt', 'w') as f:
-        for smi in smiles:
-            f.write(smi + '\n')
-
-    default_dir = get_download_dir()
-    vocab_file = '{}/jtnn/{}.txt'.format(default_dir, 'vocab')
-    zip_file_path = '{}/jtnn.zip'.format(default_dir)
-    download(_get_dgl_url('dataset/jtnn.zip'), path=zip_file_path, overwrite=False)
-    extract_archive(zip_file_path, '{}/jtnn'.format(default_dir))
-
-    with open(vocab_file, 'r') as f:
-        vocab = Vocab([x.strip("\r\n ") for x in f])
-    dataset = JTVAEDataset('data.txt', vocab)
-    assert len(dataset) == 2
-    assert set(dataset[0].keys()) == {'cand_graphs', 'mol_graph', 'mol_tree',
-                                      'stereo_cand_graphs', 'stereo_cand_label',
-                                      'tree_mess_src_e', 'tree_mess_tgt_e',
-                                      'tree_mess_tgt_n', 'wid'}
-    dataset.training = False
-    assert set(dataset[0].keys()) == {'mol_graph', 'mol_tree', 'wid'}
-
-    dataset.training = True
-    collate_fn = JTVAECollator(training=True)
-    loader = DataLoader(dataset,
-                        batch_size=2,
-                        collate_fn=collate_fn)
-    for _, batch_data in enumerate(loader):
-        assert set(batch_data.keys()) == {'cand_batch_idx', 'cand_graph_batch', 'mol_graph_batch',
-                                          'mol_trees', 'stereo_cand_batch_idx',
-                                          'stereo_cand_graph_batch', 'stereo_cand_labels',
-                                          'stereo_cand_lengths', 'tree_mess_src_e',
-                                          'tree_mess_tgt_e', 'tree_mess_tgt_n'}
-
-    dataset.training = False
-    collate_fn = JTVAECollator(training=False)
-    loader = DataLoader(dataset,
-                        batch_size=2,
-                        collate_fn=collate_fn)
-    for _, batch_data in enumerate(loader):
-        assert set(batch_data.keys()) == {'mol_graph_batch', 'mol_trees'}
-
-    remove_file('data.txt')
-    remove_file(zip_file_path)
-    remove_dir(default_dir + '/jtnn')
+    vocab = JTVAEVocab()
+    dataset = JTVAEZINC('train', vocab)
 
 if __name__ == '__main__':
-    test_alchemy()
-    test_pdbbind()
-    test_pubchem_aromaticity()
-    test_tox21()
-    test_esol()
-    test_freesolv()
-    test_lipophilicity()
-    test_astrazeneca_chembl_solubility()
-    test_wln_reaction()
+    # test_alchemy()
+    # test_pdbbind()
+    # test_pubchem_aromaticity()
+    # test_tox21()
+    # test_esol()
+    # test_freesolv()
+    # test_lipophilicity()
+    # test_astrazeneca_chembl_solubility()
+    # test_wln_reaction()
     test_jtvae()
diff --git a/tests/model/test_generative_models.py b/tests/model/test_generative_models.py
index 0173e12b..1cacdf9c 100644
--- a/tests/model/test_generative_models.py
+++ b/tests/model/test_generative_models.py
@@ -3,11 +3,10 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import torch
-
 from rdkit import Chem
 
-from dgllife.model import DGMG, DGLJTNNVAE
+from dgllife.model import DGMG, JTNNVAE
+from dgllife.utils import JTVAEVocab
 
 def test_dgmg():
     model = DGMG(atom_types=['O', 'Cl', 'C', 'S', 'F', 'Br', 'N'],
@@ -33,16 +32,12 @@ def test_dgmg():
     model.eval()
     assert model(rdkit_mol=True) is not None
 
-def test_jtnn():
-    if torch.cuda.is_available():
-        device = torch.device('cuda:0')
-    else:
-        device = torch.device('cpu')
-
-    model = DGLJTNNVAE(hidden_size=1,
-                       latent_size=2,
-                       depth=1).to(device)
+def test_jtvae():
+    vocab = JTVAEVocab()
+    model = JTNNVAE(vocab,
+                    hidden_size=1,
+                    latent_size=2,
+                    depth=3)
 
 if __name__ == '__main__':
     test_dgmg()
-    test_jtnn()
diff --git a/tests/model/test_pretrain.py b/tests/model/test_pretrain.py
index 12759cd6..e192cf03 100644
--- a/tests/model/test_pretrain.py
+++ b/tests/model/test_pretrain.py
@@ -56,9 +56,7 @@ def test_jtnn():
     else:
         device = torch.device('cpu')
 
-    model = load_pretrained('JTNN_ZINC').to(device)
-
-    remove_file('JTNN_ZINC_pre_trained.pth')
+    model = load_pretrained('JTNN_ZINC_no_kl').to(device)
 
 def test_gcn_tox21():
     if torch.cuda.is_available():
diff --git a/tests/utils/test_jtvae.py b/tests/utils/test_jtvae.py
new file mode 100644
index 00000000..e7d8a46b
--- /dev/null
+++ b/tests/utils/test_jtvae.py
@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# pylint: disable= no-member, arguments-differ, invalid-name
+#
+# Test for utils in JTVAE
+
+import os
+
+def remove_file(fname):
+    if os.path.isfile(fname):
+        try:
+            os.remove(fname)
+        except OSError:
+            pass
+
+def test_vocab():
+    from dgllife.utils.jtvae import JTVAEVocab
+
+    vocab = JTVAEVocab()
+    assert vocab.get_smiles(0) == 'C1=[NH+]C=[NH+]CC1'
+    assert vocab.get_index('C1=[NH+]C=[NH+]CC1') == 0
+
+    tmp_file = 'tmp.txt'
+    with open(tmp_file, 'w') as f:
+        f.write('CCO\n')
+        f.write('C1=CC2=CC=CC=CC2=C1\n')
+        f.write('O=C(O)/C=C/C(=O)O\n')
+        f.write('N[C@@H](C)C(=O)O\n')
+    vocab = JTVAEVocab(tmp_file)
+    assert set(vocab.vocab) == \
+           set(['C=C', 'C1=CCC=C1', 'C', 'C=O', 'CN', 'C1=CC=CCC=C1', 'CO', 'CC'])
+    remove_file(tmp_file)
+
+if __name__ == '__main__':
+    test_vocab()