Skip to content

GitLab

Unverified Commit 5b67ab99 authored by Sylvain Gugger's avatar Sylvain Gugger Committed by GitHub
Browse files

Fix import in Accelerate for find_exec_bs (#21501)

parent eb1771ef
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
src/transformers/trainer_utils.py
View file @ 5b67ab99
...@@ -646,9 +646,9 @@ def find_executable_batch_size( ...@@ -646,9 +646,9 @@ def find_executable_batch_size(
if auto_find_batch_size: if auto_find_batch_size:
requires_backends(find_executable_batch_size, "accelerate") 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) return functools.partial(function, batch_size=starting_batch_size)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment