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)