how to get the batch.ligand_element_batch

[545487677]

Hi, thank you for sharing such a good work. However, I am a little confused about how can I get batch.ligand_element_batch in the def train(it): model.train() optimizer.zero_grad() for _ in range(config.train.n_acc_batch): batch = next(train_iterator).to(args.device)

        results = model.get_diffusion_loss(
            ligand_pos=batch.ligand_pos, #
            ligand_v=batch.ligand_atom_feature_full,
            batch_ligand=batch.ligand_element_batch
        )

Can you tell where can I find the processing operation of the ligand element batch? Thank you