Combining deterministic rules with neural network

Hello everyone,

I want to set up a classification model that, in a condensed version, has 4 input features [x_1, x_2, x_3, x_4] and predicts 4 classes {y_1, y_2, y_3, y_4}. Now there exits a rule that says:
if x_1 < 0,6 then the the class can only be y_1 or y_2
else the class can only be y_3 or y_4.

They way i thought of, would be to train separate models for the 2 cases and then combine them with a simple if condition. However, given that this solution does not result in a pytorch nn.module at inference it would be super slow, considering the if condition would only be used on a single input.
This could be circumvented by collecting multiple inputs with the same result for the if condition and then feeding them to the model as a vector. However, considering in reality i will have multiple if conditions and i would like to create a solution that is dynamic with respect the these conditions (so i could easily add, remove, adapt them without changing much of the models code) and further more i am working with image data so i have to rearrange the predictions for each pixel afterwards into the same order, it seems rather complicated to keep track of all the positions of the singular inputs.

So long story short - my question is: Is there a way to model a deterministic decision as a pytorch layer? Or is there a way to build a pytoch neural network representing my desired model, containing deterministic decisions?

Best wishes,