ls_mlkit.util.offload.graph_hook module¶

class ls_mlkit.util.offload.graph_hook.OffloadSavedTensorHook[source]¶

Bases: object

mode = 'release'¶

offload_device = 'cpu'¶

static pack(x: Tensor)[source]¶

static unpack(packed)[source]¶

ls_mlkit.util.offload.graph_hook.offload_condition(x: Tensor) → bool[source]¶: The offload conditions are: the tensor is on cuda, and the number of bytes occupied by the tensor is equal to the theoretical number of bytes