code: def train():
      dataset.training = True
      dataloader = DataLoader(
          dataset,
          batch_size=batch_size,
          shuffle=True,
          num_workers=4,
          collate_fn=JTNNCollator(vocab, True),
          drop_last=True,
          worker_init_fn=worker_init_fn,
      )
  
      for epoch in range(MAX_EPOCH):
          word_acc, topo_acc, assm_acc, steo_acc = 0, 0, 0, 0
  
          for it, batch in enumerate(tqdm.tqdm(dataloader)):
              model.zero_grad()
              try:
                  loss, kl_div, wacc, tacc, sacc, dacc = model(
                      batch, 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) % PRINT_ITER == 0:
                  word_acc = word_acc / PRINT_ITER * 100
                  topo_acc = topo_acc / PRINT_ITER * 100
                  assm_acc = assm_acc / PRINT_ITER * 100
                  steo_acc = steo_acc / 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(),
                      opts.save_path
                      + "/model.iter-%d-%d" % (epoch, it + 1),
                  )
  
          scheduler.step()
          print("learning rate: %.6f" % scheduler.get_lr()[0])
          torch.save(
              model.state_dict(),
              opts.save_path + "/model.iter-" + str(epoch),
          )
  
file_path: dgl/examples/pytorch/jtnn/vaetrain_dgl.py
imports: import math
  import random
  import sys
  from collections import deque
  from optparse import OptionParser
  import rdkit
  import torch
  import torch.nn
  import torch.optim
  import torch.optim.lr_scheduler
  import tqdm
  from jtnn import *
  from torch.utils.data import DataLoader
label_desc: Use image node_id 2 for calling a global function with example usage: train() without return types