I would suggest that you run it with the profiling executor a few times with inputs that cover the different Tensor dimensions you expect to use, and then add a pass to remove Bailout nodes. This should give you a Graph with shape information to maximum generality.