How to use a layer with gradient but without weight adjustment?
Is it possible to mark part of the forward pass to only backpropagate the gradient but not to adjust weights?
In the following example code I have a Module that uses only one layer (one set of parameters) but it is used twice in the forward step. During the optimization I would expect the weights to be adjusted twice as well. If I want to only adjust the weights for one of the layer usages, what can I do?
import torch
    
class ExampleModel(torch.nn.Module):
    
    def __init__(self, dim) -&gt; None:
        super(ExampleModel, self).__init__()
        self.linear = torch.nn.Linear(dim, dim)
    
    def forward(self, x):
        out1 = self.linear(x)  # backprop gradients and adjust weights here
        out2 = self.linear(out1)  # only backprop gradients here
        return out2
    
    
# Random input output data for this example
N, D = 64, 100
x = torch.randn(N, D)
y = torch.randn(N, D)
    
model = ExampleModel(D)
    
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.Adam(model.parameters())
    
y_pred = model(x)
loss = criterion(y_pred, y)
    
optimizer.zero_grad()
loss.backward()
optimizer.step()

The following will not work since with torch.no_grad() no gradient at all is backpropagated:
def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    with torch.no_grad():
        out2 = self.linear(out1)  # only backprop gradients here
    return out2

I can not simply exclude the parameters from the optimization since they need to be optimized in the first part (i.e. out1 = self.linear(x)).
For the same reason I can also not set a learning rate of 0 for these parameters.
What else can I do to achieve this?

One way to do it is to use requires_grad_ to temporarily disable gradients on the layer's parameters:
def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    self.linear.requires_grad_(False)
    out2 = self.linear(out1)  # only backprop gradients here
    self.linear.requires_grad_(True)
    return out2

This still lets gradients flow through the activations; it merely stops them from reaching the parameters.
You could also consider manipulating the weight tensors manually and calling .detach():
import torch.nn.functional as F
def forward(self, x):
    out1 = self.linear(x)
    out2 = F.linear(out1, self.linear.weight.detach(), self.linear.bias.detach())
    return out2

