ls_mlkit.util.offload.gradient_offload_v2 module¶

class ls_mlkit.util.offload.gradient_offload_v2.GradientOffloadHookContext(model: Module, enable: bool, *args, **kwargs)[source]¶

Bases: object

static get_offload_grad_hook(offload_attr_name='offloaded_grad', offload_device='cpu', init_iters='init_iters', *arg, **kwargs)[source]¶

ls_mlkit.util.offload.gradient_offload_v2.get_offload_grad_hook(offload_attr_name='offloaded_grad', offload_device='cpu', init_iters='init_iters', *arg, **kwargs)[source]¶