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()