code: def move_to_cuda(mol_batch):
      for i in range(len(mol_batch["mol_trees"])):
          mol_batch["mol_trees"][i].graph = cuda(
              mol_batch["mol_trees"][i].graph
          )
  
      mol_batch["mol_graph_batch"] = cuda(mol_batch["mol_graph_batch"])
      if "cand_graph_batch" in mol_batch:
          mol_batch["cand_graph_batch"] = cuda(
              mol_batch["cand_graph_batch"]
          )
      if mol_batch.get("stereo_cand_graph_batch") is not None:
          mol_batch["stereo_cand_graph_batch"] = cuda(
              mol_batch["stereo_cand_graph_batch"]
          )
  
imports: import copy
  import rdkit.Chem
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import batch, unbatch
  from .chemutils import attach_mols_nx, copy_edit_mol, decode_stereo, enum_assemble_nx, set_atommap
  from .jtmpn import DGLJTMPN, mol2dgl_single
  from .jtnn_dec import DGLJTNNDecoder
  from .jtnn_enc import DGLJTNNEncoder
  from .mpn import DGLMPN, mol2dgl_single
  from .nnutils import cuda
label_desc: Use image node_id 2 for calling the DGLJTNNVAE obj's underlying member method code with example usage: obj.move_to_cuda(mol_batch) without return types