Function _get_group_onload_device

src/diffusers/hooks/group_offloading.py:1021–1025

(module: torch.nn.Module)

def _get_group_onload_device(module: torch.nn.Module) -> torch.device:
    top_level_group_offload_hook = _get_top_level_group_offload_hook(module)
    if top_level_group_offload_hook is not None:
        return top_level_group_offload_hook.config.onload_device
    raise ValueError("Group offloading is not enabled for the provided module.")


def _compute_group_hash(group_id):