ec54a0b75d93cab5ad407920f2bbfd2f33251cc9,pytext/trainers/trainer.py,Trainer,train,#Trainer#Any#Any#Any#Any#Any#Any#Any#Any#,62

Before Change

            sys.stdout.flush()

        if rank == 0:
            model.load_state_dict(torch.load(best_model_path))
        return model, best_metric

    def _run_epoch(
        self,

After Change

                    if cuda_utils.CUDA_ENABLED:
                        for key, state in model_state.items():
                            model_state[key] = state.cpu()
                    best_model_state = model_state

            if self.config.early_stop_after > 0 and (
                epoch - last_best_epoch == self.config.early_stop_after
            ):
                print(
                    f"Rank {rank} worker: Eval metric hasn"t changed for "
                    + f"{self.config.early_stop_after} epochs. Stopping now."
                )
                break
            sys.stdout.flush()

        if rank == 0:
            if cuda_utils.CUDA_ENABLED:
                for key, state in best_model_state.items():
                    best_model_state[key] = state.cuda()
            model.load_state_dict(best_model_state)
        return model, best_metric

    def _run_epoch(

Instances