[deepspeed] fix load_best_model test (#17550)

2022-06-03 11:19:03 -07:00
parent 72f5b94984
commit 26e5e129b4
1 changed files with 52 additions and 52 deletions
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -752,6 +752,8 @@ class TrainerIntegrationDeepSpeed(TrainerIntegrationDeepSpeedWithCustomConfig, T
        # must use this setting to get the reload path exercised
        ds_config_dict["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] = True

+        with mockenv_context(**self.dist_env_1_gpu):
+
            tokenizer = T5Tokenizer.from_pretrained(T5_TINY)
            model = T5ForConditionalGeneration.from_pretrained(T5_TINY)

@@ -804,8 +806,6 @@ class TrainerIntegrationDeepSpeed(TrainerIntegrationDeepSpeedWithCustomConfig, T
                "deepspeed": ds_config_dict,
            }

-        with mockenv_context(**self.dist_env_1_gpu):
-
            training_args = TrainingArguments(output_dir, **args_dict)

            trainer = Trainer(