pytorch/torch/amp/__init__.py

10 lines
181 B
Python

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