ff68a9ef501e7286501dba1719024dfaaab4b473,tests/test_train.py,TestLoadCheckpoint,test_load_partial_checkpoint,#TestLoadCheckpoint#,41

Before Change

        with contextlib.redirect_stdout(StringIO()):
            trainer = mock_trainer(2, 200, False)
            loader = mock_loader(150)
            epoch, ds = train.load_checkpoint(MagicMock(), trainer, loader)
            self.assertEqual(epoch, 2)
            self.assertEqual(next(ds), 50)

After Change

    def test_load_partial_checkpoint(self):
        with contextlib.redirect_stdout(StringIO()):
            trainer, epoch_itr = get_trainer_and_epoch_itr(2, 150, 200, 50)

            train.load_checkpoint(MagicMock(), trainer, epoch_itr)
            self.assertEqual(epoch_itr.epoch, 2)
            self.assertEqual(epoch_itr.iterations_in_epoch, 50)

            itr = epoch_itr.next_epoch_itr(shuffle=False)
            self.assertEqual(epoch_itr.epoch, 2)
            self.assertEqual(epoch_itr.iterations_in_epoch, 50)

            self.assertEqual(next(itr)["net_input"]["src_tokens"][0].item(), 50)

Instances