File size: 667 Bytes
fa2d87b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | from __future__ import annotations
def enable_h3_cpu_offload(
components_manager, *, memory_reserve_margin: str = "64GB"
) -> dict[str, str]:
device = "cuda"
components_manager.enable_auto_cpu_offload(
device=device,
memory_reserve_margin=memory_reserve_margin,
)
return {
"mode": "components_manager_auto_cpu_offload",
"device": device,
"memory_reserve_margin": memory_reserve_margin,
}
def component_device(module) -> str:
tensor = next(module.parameters(), None)
if tensor is None:
tensor = next(module.buffers(), None)
return "none" if tensor is None else str(tensor.device)
|