Fix import in Accelerate for find_exec_bs (#21501)

huggingface · Feb 9, 2023 · 865bb4e · 865bb4e
1 parent 02c3f41
commit 865bb4e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/trainer_utils.py b/src/transformers/trainer_utils.py
@@ -646,9 +646,9 @@ def find_executable_batch_size(
 
     if auto_find_batch_size:
         requires_backends(find_executable_batch_size, "accelerate")
-        import accelerate.memory_utils as mem_utils
+        from accelerate.utils import find_executable_batch_size as accelerate_find_executable_batch_size
 
-        return mem_utils.find_executable_batch_size(function=function, starting_batch_size=starting_batch_size)
+        return accelerate_find_executable_batch_size(function=function, starting_batch_size=starting_batch_size)
 
     return functools.partial(function, batch_size=starting_batch_size)