fix(utils): extraction of mixed_precision option for deepspeed (#197)

CarperAI · Jan 16, 2023 · 830f490 · 830f490
1 parent c25f598
commit 830f490
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/trlx/utils/__init__.py b/trlx/utils/__init__.py
@@ -55,10 +55,9 @@ def get_distributed_config(accelerator: Accelerator):
     Return accelerator distributed config
     """
 
-    accelerate_config = accelerator.state
     dist_config = {
-        "mixed_precision": accelerate_config.mixed_precision,
-        "num_gpus": accelerate_config.num_processes,
+        "mixed_precision": accelerator.mixed_precision,
+        "num_gpus": accelerator.num_processes,
     }
 
     if accelerator.state.deepspeed_plugin is not None: