def train_step(
        self, sample, model, criterion, optimizer, update_num, ignore_grad=False
    ):
        """
        Do forward and backward, and return the loss as computed by *criterion*
        for the given *model* and *sample*.

        Args:
            sample (dict): the mini-batch. The format is defined by the
                :class:`~fairseq.data.FairseqDataset`.
            model (~fairseq.models.BaseFairseqModel): the model
            criterion (~fairseq.criterions.FairseqCriterion): the criterion
            optimizer (~fairseq.optim.FairseqOptimizer): the optimizer
            update_num (int): the current update
            ignore_grad (bool): multiply loss by 0 if this is set to True

        Returns:
            tuple:
                - the loss
                - the sample size, which is used as the denominator for the
                  gradient
                - logging outputs to display while training
        """
        model.train()
        model.set_num_updates(update_num)

        agg_loss, agg_sample_size, agg_logging_output = 0., 0., {}

        with torch.autograd.profiler.record_function("forward"):
            loss, sample_size, logging_output = criterion(model, sample['gpt'], loss_name='gpt')
        if ignore_grad:
            loss *= 0
        with torch.autograd.profiler.record_function("backward"):
            optimizer.backward(loss)
        
        agg_loss += loss.detach().item()
        agg_sample_size += sample_size
        agg_logging_output.update(logging_output)
        
        if 'laion' in sample:
            with torch.autograd.profiler.record_function("forward"):
                loss, sample_size, logging_output = criterion(model, sample['laion'], loss_name='laion')
            if ignore_grad:
                loss *= 0
            with torch.autograd.profiler.record_function("backward"):
                optimizer.backward(loss)
            
            agg_loss += loss.detach().item()
            agg_sample_size += sample_size
            for key, value in logging_output.items():
                if key not in agg_logging_output:
                    agg_logging_output[key] = value
                else:
                    agg_logging_output[key] += value

        if 'wild' in sample:
            with torch.autograd.profiler.record_function("forward"):
                loss, sample_size, logging_output = criterion(model, sample['wild'], loss_name='wild')
            if ignore_grad:
                loss *= 0
            with torch.autograd.profiler.record_function("backward"):
                optimizer.backward(loss)
            
            agg_loss += loss.detach().item()
            agg_sample_size += sample_size
            for key, value in logging_output.items():
                if key not in agg_logging_output:
                    agg_logging_output[key] = value
                else:
                    agg_logging_output[key] += value

        return agg_loss, agg_sample_size, agg_logging_output

        Do forward and backward, and return the loss as computed by *criterion*
        for the given *model* and *sample*.

        Args:
            sample (dict): the mini-batch. The format is defined by the
                :class:`~fairseq.data.FairseqDataset`.
            model (~fairseq.models.BaseFairseqModel): the model
            criterion (~fairseq.criterions.FairseqCriterion): the criterion
            optimizer (~fairseq.optim.FairseqOptimizer): the optimizer
            update_num (int): the current update
            ignore_grad (bool): multiply loss by 0 if this is set to True

        Returns:
            tuple:
                - the loss
                - the sample size, which is used as the denominator for the
                  gradient
                - logging outputs to display while training
        
kosmos-2/torchscale/examples/fairseq/tasks/vl_gpt_base.py