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

* 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

View File

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