oneflow.nn.graph.graph_config.GraphConfig.allow_fuse_model_update_ops

GraphConfig.allow_fuse_model_update_ops(mode: bool = True)

If set to true, try to fuse cast + scale + l1_l2_regularize_gradient + model_update to one op to improve performance.

For example:

import oneflow as flow

class Graph(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.linear = flow.nn.Linear(3, 8, False)
        self.config.allow_fuse_model_update_ops(True)
    def build(self, x):
        return self.linear(x)

graph = Graph()
Parameters

mode (bool, optional) – The default value is True.