def METHOD_NAME(
    self, sample, model, criterion, optimizer, update_num, ignore_grad=False
):
    model.train()
    model.set_num_updates(update_num)
    # task ids
    if self.task_type == "lang":
        encoder_task_ids = sample["src_lang_ids"]
        decoder_task_ids = sample["tgt_lang_ids"]
    elif self.task_type == "domain":
        encoder_task_ids = sample["domain_ids"]
        decoder_task_ids = sample["domain_ids"]
    model.encoder.set_task_ids(encoder_task_ids)
    model.decoder.set_task_ids(decoder_task_ids)
    with torch.autograd.profiler.record_function("forward"):
        with torch.cuda.amp.autocast(enabled=(isinstance(optimizer, AMPOptimizer))):
            loss, sample_size, logging_output = criterion(model, sample)
            # KL loss
            if self.args.encoder_attn_head_select:
                sample_sizes = self.get_sample_sizes(sample, encoder_task_ids, self.encoder_tasks)
                loss += self.kl_loss(
                    model.encoder.attn_head_selector.head_samples,
                    sample_sizes,
                    self.encoder_head_prior
                )
            if self.args.decoder_self_attn_head_select:
                sample_sizes = self.get_sample_sizes(sample, decoder_task_ids, self.decoder_tasks)
                loss += self.kl_loss(
                    model.decoder.self_attn_head_selector.head_samples,
                    sample_sizes,
                    self.decoder_head_prior
                )
            if self.args.dec_enc_attn_head_select:
                sample_sizes = self.get_sample_sizes(sample, decoder_task_ids, self.decoder_tasks)
                loss += self.kl_loss(
                    model.decoder.enc_attn_head_selector.head_sampes,
                    sample_sizes,
                    self.decoder_head_prior
                )
    if ignore_grad:
        loss *= 0
    with torch.autograd.profiler.record_function("backward"):
        optimizer.backward(loss)
    return loss, sample_size, logging_output