The ConvertLayout pass works well but unfortunately leaves layout_transform ops in the graph, which can be a real pain to remove. I used pattern rewrite but I’d like a more generic solution.
More specifically, I’d like to convert a full model (including inputs and outputs) from one layout to another, say NCHW to NHWC. The resulting graph must be clean without layout_transform, just the original ops with tensors in the final layout and weight tensors in corresponding layout eg for convolutions.
Is there a generic way to do that?