acc_loss = 0
count = 0
bar = tqdm(train_loader, total = len(train_loader))
for i, ((old_node_token_ids, old_node_type_ids), (new_node_token_ids, new_node_type_ids), old_token_ids, batch_joint_graph_dgl, new_graph_dgl, dual_graph_dgl, changed_node_ids, batch_diff_token_ids, metadata) in enumerate(bar):
batch_diff_token_ids = batch_diff_token_ids.long()
tgt_input = batch_diff_token_ids[:, :-1]
tgt_label = batch_diff_token_ids[:, 1:]
old_token_ids = torch.split(old_token_ids, metadata['old_token_num'])
pad_old_token_ids = list(map(lambda x: x.numpy().tolist(), old_token_ids))
pad_old_token_ids = _pad_batch_2D(pad_old_token_ids)
pad_old_token_ids = torch.Tensor(pad_old_token_ids).long()
old_token_embed = comment_embedding(pad_old_token_ids)
memory = encoder(old_token_embed, src_key_padding_mask = pad_old_token_ids == 0)
initial_tar_embeds = comment_embedding(tgt_input)
tgt_mask = generate_square_subsequent_mask(initial_tar_embeds.shape[1])#.to(old_token_mask.device)
tar_embedding = decoder(initial_tar_embeds, memory, tgt_mask = tgt_mask, memory_key_padding_mask = pad_old_token_ids == 0, tgt_key_padding_mask = tgt_input == 0)
tar_output_gen = out_fc(tar_embedding).transpose(1, 2)
loss = nn.CrossEntropyLoss(ignore_index = 0)(tar_output_gen, tgt_label)
loss.backward()
opt.step()
count += 1
acc_loss += loss.cpu().item()
val_loss = eval(model, valid_loader)
print('train_loss', acc_loss / len(train_loader), 'val loss', val_loss)