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


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