frozenleaves/accelerate
1
0
Fork 0
mirror of https://github.com/huggingface/accelerate.git synced 2025-10-20 18:13:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Use with in Accelerator.autocast()instead of __enter__() and __exit__() for more elegant style. (#3767)

Browse Source 
* refactor: Use ` with`  in `Accelerator.autocast()`instead of  `__enter__()` and `__exit__()`for more elegant style.

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Walker
2025-09-18 21:27:12 +08:00
committed by GitHub
parent 3db9fb6991
commit 8031e24e84
1 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/accelerate/accelerator.py
View File

@ -21,7 +21,6 @@ import math
import os import os
import re import re
import shutil import shutil
import sys
import warnings import warnings
from collections import OrderedDict from collections import OrderedDict
from contextlib import contextmanager from contextlib import contextmanager
@ -4069,10 +4068,8 @@ class Accelerator:
if autocast_handler is None: if autocast_handler is None:
autocast_handler = self.autocast_handler autocast_handler = self.autocast_handler
autocast_context = get_mixed_precision_context_manager(self.native_amp, autocast_handler) autocast_context = get_mixed_precision_context_manager(self.native_amp, autocast_handler)
autocast_context.__enter__() with autocast_context:
# TODO: should the `yield` be in a try/finally block? yield
yield
autocast_context.__exit__(*sys.exc_info())
@contextmanager @contextmanager
def profile(self, profile_handler: ProfileKwargs | None = None): def profile(self, profile_handler: ProfileKwargs | None = None):