Is there a way that embeds a if-else structure into a network

HI, all, I want to know if there is a way that can embed a if-else structure into a network. For example, in a feature map, for the F(u,v) value in (u,v) location, I want to replace this value with the mean value of its smoothest neighbors. We have to judge which neigbor is smooth, So we will use a if-else structure in the network.

However, I think it cannot be trained if use if-else structure directly. So if there is a way to solve it.

PS: such a operator cannot be expressed in a function so we canno make it a torch.autograd.Funcion.

You can just create a simple function for the your purpose and call that function directly. PyTorch would automatically handle the backprop step for you.
For the if-else part, you can add if-else statements in your model and pytorch would train them for you.