code: 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
      self.embedding = embedding
  
      if embedding is None:
          self.embedding = nn.Embedding(self.vocab_size, hidden_size)
      else:
          self.embedding = embedding
  
      # GRU Weights
      self.W_z = nn.Linear(2 * hidden_size, hidden_size)
      self.U_r = nn.Linear(hidden_size, hidden_size, bias=False)
      self.W_r = nn.Linear(hidden_size, hidden_size)
      self.W_h = nn.Linear(2 * hidden_size, hidden_size)
  
      # Word Prediction Weights
      self.W = nn.Linear(hidden_size + latent_size, hidden_size)
  
      # Stop Prediction Weights
      self.U = nn.Linear(hidden_size + latent_size, hidden_size)
      self.U_i = nn.Linear(2 * hidden_size, hidden_size)
  
      # Output Weights
      self.W_o = nn.Linear(hidden_size, self.vocab_size)
      self.U_o = nn.Linear(hidden_size, 1)
  
      # Loss Functions
      self.pred_loss = nn.CrossEntropyLoss(size_average=False)
      self.stop_loss = nn.BCEWithLogitsLoss(size_average=False)
  
file_path: DIG/dig/ggraph/method/JTVAE/fast_jtnn/jtnn_dec.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from .mol_tree import MolTree, MolTreeNode
  from .chemutils import enum_assemble, set_atommap
  from .nnutils import create_var, GRU
label_desc: Use image node_id 1 to create a new JTNNDecoder object from inherited base classes: nn with example: obj = JTNNDecoder(vocab, hidden_size, latent_size, embedding)