venv/lib/python3.12/site-packages/torch/amp/__init__.py

from .autocast_mode import (
    _enter_autocast,
    _exit_autocast,
    autocast,
    custom_bwd,
    custom_fwd,
    is_autocast_available,
)
from .grad_scaler import GradScaler