code: def train_model(
      model, src_words_idx, tgt_words_idx, tgt_seq_len, src_seq_len
  ):
      optimizer = optim.Adam(model.parameters(), lr=5e-3)
      mask = seq_len_to_mask(tgt_seq_len).eq(0)
      target = tgt_words_idx.masked_fill(mask, -100)
  
      for i in range(50):
          optimizer.zero_grad()
          pred = model(src_words_idx, tgt_words_idx, src_seq_len)[
              "pred"
          ]  # bsz x max_len x vocab_size
          loss = F.cross_entropy(pred.transpose(1, 2), target)
          loss.backward()
          optimizer.step()
  
      right_count = (
          pred.argmax(dim=-1).eq(target).masked_fill(mask, 1).sum()
      )
      return right_count
  
file_path: fastNLP/tests/models/torch/test_seq2seq_model.py
imports: import pytest
  from fastNLP.envs.imports import _NEED_IMPORT_TORCH
  from fastNLP import seq_len_to_mask
label_desc: Use image node_id 2 for calling a global function with example usage: train_model(model, src_words_idx, tgt_words_idx, tgt_seq_len, src_seq_len) and returns: right_count