oneflow.nn.graph.block_config.BlockConfig.activation_checkpointing

property BlockConfig.activation_checkpointing

Set/Get whether do activation checkpointing in this nn.Module.

For example:

import oneflow as flow

class Graph(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.linear1 = flow.nn.Linear(3, 5, False)
        self.linear2 = flow.nn.Linear(5, 8, False)
        self.linear1.config.activation_checkpointing = True
        self.linear2.config.activation_checkpointing = True

    def build(self, x):
        y_pred = self.linear1(x)
        y_pred = self.linear2(y_pred)
        return y_pred

graph = Graph()