Untitled

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)