def METHOD_NAME(self, graph_def, arrays):
    meta_graph = export_meta_graph(graph_def=graph_def)
    fetch_collection = meta_graph_pb2.CollectionDef()
    for name in arrays:
        fetch_collection.node_list.value.append(name)
    meta_graph.collection_def["train_op"].CopyFrom(fetch_collection)
    # Initialize RewriterConfig with everything disabled except function
    # inlining.
    config = tf.compat.v1.ConfigProto()
    rewrite_options = config.graph_options.rewrite_options
    rewrite_options.optimizers.append("function")
    return tf_optimizer.OptimizeGraph(config, meta_graph)