ls_mlkit.util.offload.context module

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

Bases: object