f4f19f8ae940c66c1f8af4c3aaacb9e0b6cbbbf9,pytext/task/new_task.py,NewTaskTrainer,_run_epoch,#NewTaskTrainer#Any#Any#Any#Any#Any#Any#Any#Any#Any#,29

Before Change

        Our run_epoch is a bit different, because we"re wrapping the model forward
        call with model.train_batch, which arranges tensors and gets loss, etc.
        print(f"Rank {rank} worker: Running epoch //{epoch} for {stage}")
        timer = time_utils.StageTimer()
        report_metric = stage != Stage.TRAIN or self.config.report_train_metrics

        for batch_id, batch in enumerate(batches):
            pre_batch()

After Change

        for batch_id, batch in enumerate(batches):
            pre_batch()
            with time_utils.time("model.train_batch"):
                loss, metric_data = model.train_batch(batch)
            with time_utils.time("backprop"):
                backprop(loss)
            if report_metric:
                with time_utils.time("add metrics"):
                    metric_reporter.add_batch_stats(
                        batch_id, *metric_data, **metric_reporter.batch_context(batch)
                    )

        metrics = None
        if report_metric:
            with time_utils.time("report metrics"):
                metrics = metric_reporter.report_metric(
                    stage, epoch, print_to_channels=(rank == 0)
                )
        else:
            metric_reporter._reset()

        return metrics

Instances