ls_mlkit.util.offload.gradient_offload module

class ls_mlkit.util.offload.gradient_offload.GradientOffloadHookContext(model: Module, record_dict: dict, enable: bool = True, *args, **kwargs)

Bases: object

get_record_gradient_hook(model, record_dict)

register_gradient_hook()

ls_mlkit.util.offload.gradient_offload.get_record_gradient_hook(self, model, record_dict)