__init__.py 314 Bytes
Newer Older
Aryan's avatar
Aryan committed
1
2
3
4
from ..utils import is_torch_available


if is_torch_available():
5
    from .hooks import HookRegistry, ModelHook
Aryan's avatar
Aryan committed
6
    from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
7
    from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast