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?