ls_mlkit.util.offload package

Submodules

• ls_mlkit.util.offload.context module
• ls_mlkit.util.offload.forward_backward_offload module
• ls_mlkit.util.offload.forward_hook module
• ls_mlkit.util.offload.gradient_offload module
• ls_mlkit.util.offload.gradient_offload_v2 module
• ls_mlkit.util.offload.graph_hook module
• ls_mlkit.util.offload.model_offload module
• ls_mlkit.util.offload.resource_monitor module
• ls_mlkit.util.offload.saved_tensor_offload module
• ls_mlkit.util.offload.split module

Module contents

class ls_mlkit.util.offload.OffloadContext(model: Module, named_grads: dict, num_block=2, no_split_module_classes=None, enable_gradient_offload=True, enable_model_offload=True)

Bases: object