[deepspeed] fix load_best_model test (#17550)

26e5e129 · Stas Bekman · GitHub · 72f5b949 · 26e5e129
Unverified Commit 26e5e129 authored Jun 03, 2022 by Stas Bekman Committed by GitHub Jun 03, 2022
Show whitespace changes
Inline Side-by-side

Showing with 51 additions and 51 deletions

tests/deepspeed/test_deepspeed.py tests/deepspeed/test_deepspeed.py +51 -51

No files found.
--- 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(