gpt2_trial

determined-ai · Oct 31, 2024 · 41a2cb3 · 41a2cb3
1 parent 4502276
commit 41a2cb3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/deepspeed/gpt_neox/gpt2_trial.py b/examples/deepspeed/gpt_neox/gpt2_trial.py
@@ -23,6 +23,7 @@ class GPT2Trial(det_ds.DeepSpeedTrial):
     def __init__(
         self, context: det_ds.DeepSpeedTrialContext, hparams: dict, trial_seed: int
     ) -> None:
+        self.context = context
         self.exp_config = context.get_experiment_config()
         self.args = attrdict.AttrMap(hparams)