How do I make a custom layer that will be used only during prediction stage?

During the prediction of a trained neural network, I want to add new custom final layer and get its output. How do I do that?

I created a custom layer with and added it as the last layer of my network. But, that layer is still called during the forward pass during training. How do I do this right?

During the forward you can use that condition:
if do_this else dont is a default flag in nn.Module.

@JuanFMontesinos Thanks. I misunderstood what was for. Yes, that seems to work. But, if I’m training like this it won’t work :

            output = model(x_train)
            loss_value = loss(output, y_train)

            # backprop

What do you mean by “like this”
If you set model.eval() it won’t call it, else it will.