Does adding a forward hook to a layer of ensure that the gradient of the loss calculated using the layer's output will be calculated automatically?
I have a model
class NewModel(nn.Module):
    def __init__(self,output_layer,*args):
        self.output_layer = output_layer
        super().__init__(*args)

        self.output_layer = output_layer
        self.selected_out = None
        #PRETRAINED MODEL
        self.pretrained = models.resnet18(pretrained=True)
    
        #TAKING OUTPUT FROM AN INTERMEDIATE LAYER

        #self._layers = []
        for l in list(self.pretrained._modules.keys()):
            #self._layers.append(l)
            if l == self.output_layer:
                handle = getattr(self.pretrained,l).register_forward_hook(self.hook)
   
    def hook(self,module, input,output):
        self.selected_out = output

    def forward(self, x):
        return x = self.pretrained(x)

I have two target outputs, one which is same as any label of an image and the second one is the same dimensions as the output obtained from self.output_layer, called target_feature
out = model(img)
layerout = model.selected_out

Now, if I want to calculate the loss of layerout with the target feature map, can it be done like the line written below?
loss = criterion(y_true, out) + feature_criterion(layerout, target_feature)

Or do I need to add backward_hooks?
In this Kaggle notebook
https://www.kaggle.com/sironghuang/understanding-pytorch-hooks
it is written that loss.backward() cannot be used when using backward_hooks.
Quoting the author
# backprop once to get the backward hook results
out.backward(torch.tensor([1,1],dtype=torch.float),retain_graph=True)
#! loss.backward(retain_graph=True)  # doesn't work with backward hooks, 
#! since it's not a network layer but an aggregated result from the outputs of last layer vs target 

Then how can be gradient be calculated based on the loss function?

If I understand you correctly, you want to get two outputs from your model, calculate two losses, then combine them and backpropagate. I imagine you come from Tensorflow &amp; Keras from the way you tried implementing it. In Pytorch, it's actually fairly straight foward, you can do this very easily because of its purely functional aspect.
This is just an example:
class NewModel(nn.Module):
    def __init__(self, output_layer, *args):
        super(MyModel, self).__init__()
        
        self.pretrained = models.resnet18(pretrained=True)
        self.output_layer = output_layer
        
    def forward(self, x):
        out = self.pretrained(x)
        features = self.output_layer(out)
        return out, features

On inference, you will get two results per call:
&gt;&gt;&gt; m = NewModel(nn.Linear(1000, 10))
&gt;&gt;&gt; x = torch.rand(16, 3, 224, 224)
&gt;&gt;&gt; y_pred, y_feature = m(x)

Call you loss functions:
&gt;&gt;&gt; loss = criterion(y_pred, y_true) + feature_criterion(y_feature, target_feature)

Then, backpropagate with loss.backward().
So no need for hooks, nor complicated gradient on your .backward call!

Edit - If you wish to extract an intermediate layer output, keep the hook, that's good. And just modify the forward definition.
def forward(self, x):
    out = self.pretrained(x)
    return out, self.selected_out

For example:
&gt;&gt;&gt; m = NewModel(output_layer='layer1')
&gt;&gt;&gt; x = torch.rand(16, 3, 224, 224)
&gt;&gt;&gt; y_pred, y_feature = m(x)

&gt;&gt;&gt; y_pred.shape, y_feature.shape
(torch.Size([16, 1000]), torch.Size([16, 64, 56, 56]))

Also, what I said above about the loss stills stands. Compute your loss, then call loss.backward().
