Let’s say I have a model, which has two fully connected layers. in the forward method, I am arbitrarily changing the output tensor after first layer by squaring the tensor.
I want to understand, how will the autograd work in this case? Will this transformation affect the learned weight vectors? ideally it should. What if I don’t want to take the effect of this transformation on the gradient?
I am trying to reason along following lines :


The out tensor has requires_grad = False. So this is not going to get updated. But autograd will compute gradient of out*out w.r.t. out to propagate the gradients. right ?


The transformation out = out*out, is not having any parameters.


What if I want to apply a transformation but don’t want it to have effect on my previous layers parameters. How to achieve that ?


class NeuralNetwork(nn.Module):
     def __init__(self):
        super(NeuralNet, self).__init__()

        self.fc1 = nn.Linear(10,20)
        self.fc2 = nn.Linear(20,10) 

    def forward(self, inp):

        out = self.fc1(inp)

        # transform the out vector 
        out = out * out 

        out = self.fc2(out)

        return out