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 self.training=False 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 self.training do_this else dont self.training is a default flag in nn.Module.

@JuanFMontesinos Thanks. I misunderstood what self.training 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
            model.zero_grad()
            loss_value.backward()
            model.update_parameters(learning_rate)

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