Fix accelerate logger bug (huggingface#23650)

* fix logger bug * Update tests/mixed_int8/test_mixed_int8.py Co-authored-by: Zachary Mueller <[email protected]> * import `PartialState` --------- Co-authored-by: Zachary Mueller <[email protected]>
sheonhan · Jun 1, 2023 · 5d7ab26 · 5d7ab26
1 parent fd821a3
commit 5d7ab26
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/mixed_int8/test_mixed_int8.py b/tests/mixed_int8/test_mixed_int8.py
@@ -29,6 +29,7 @@
     pipeline,
 )
 from transformers.testing_utils import (
+    is_accelerate_available,
     is_torch_available,
     require_accelerate,
     require_bitsandbytes,
@@ -40,6 +41,13 @@
 from transformers.utils.versions import importlib_metadata
 
 
+if is_accelerate_available():
+    from accelerate import PartialState
+    from accelerate.logging import get_logger
+
+    logger = get_logger(__name__)
+    _ = PartialState()
+
 if is_torch_available():
     import torch
     import torch.nn as nn