flex_model.distributed.BaseOffloadStrategy

class flex_model.distributed.BaseOffloadStrategy(offload_fn)

Defines an offload strategy, which each device may or may not participate in. Offloading means taking the tensor and disconnecting it from any computation graph for separate downstream processing.