Let’s say I have a model, which has 2 inputs. The first input is a number, and the second input is another number, which in reality is a class.

This model is split into 2 submodels. First sub-model works on the input, and the second sub-model works on the output of the first sub-model.

The value of the first input will very greatly by the output of the second. Thus, I wish to be able to have multiple candidates of the first sub-model, and dynamically select which one to use at each step, both during training and inference, based on the value (class) of the second input.

I did not manage to achieve this. I tried using the tf.cond, the tf.switch_case and several other things, but I never managed. When I asked chat GPT it said I should be using PyTorch for this. Is there really no way to do this ?