I am trying to implement an input adaptive model (aka dynamic model), like this one: [https://arxiv.org/abs/2003.10401](dynamic model), where each input determine the structure of the model. I need to get the computational cost corresponding to that specific input for optimization. The computation cost corresponds to foo here.