Add an input to a model with torch.compile

Hello,
I have a PyTorch function that use torch.fx to add a second input to my model and performed some additional operations.
Since torch.fx is not handling control flow, I want to use torch.compile with a custom backend that will modify the graph (by adding a new input).

Any idea ?