oneflow.nn.graph.graph_config.GraphConfig.allow_fuse_cast_scale

GraphConfig.allow_fuse_cast_scale(mode: bool = True)

If set to true, try to fuse cast and scalar_mul_by_tensor to improve performance.

For example:

import oneflow as flow

def model(x):
    return flow.mul(1,flow.cast(x,flow.int8))

class Graph(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.m=model
        self.config.allow_fuse_cast_scale(True)
    def build(self, x):
        return self.m(x)

graph = Graph()
Parameters

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