Gradient computation in custom backward

Hdk: <p>How can I get gradients of ‘Li’ w.r.t. its corresponding <strong>‘Zi_unnormalized’</strong> and such that a tensor of shape same as that of feat is formed with these gradients?</p>
<pre><code class="lang-auto">    ....
    @staticmethod
    def backward(ctx, grad_output):    # grad_output size = ([batch_size, 128,4,4])
        
        input, weight, bias, output = ctx.saved_tensors    #input size = ([batch_size, 96,8,8])
      
        feat = output.clone().detach()   # output from forward with size = ([batch_size, 128,4,4])

        features = feat.view(feat.shape[0], feat.shape[1], -1) # features size = ([batch_size, 128,16])

        i = 0        
        cont = torch.tensor([0.]).to(dev)
        while i in range(0, features.shape[0]):
              for f in range(len(features[i])):
                  Zi_unnormalized = features[i][f].requires_grad_(requires_grad=True)
                  Zi = torch.nn.functional.normalize(Zi_unnormalized, dim = 0)
                  # Zj and Zk are tensors made from features[i][*] and features[other than i][*]. Zj and Zk varies for each Zi (or f)

                  Zi_Zk = torch.Tensor([0]).to(dev)
                  for k in Zk:
                      k= torch.nn.functional.normalize(k, dim = 0)
                      zi_zk = torch.exp(torch.div(torch.dot(Zi,k.T), 0.07))
                      Zi_Zk = Zi_Zk.add(zi_zk)

                  # Similarly computing Zi_Zj
                  # Li = some algebra of Zi_Zj and Zi_Zk
                  # number of 'Li' values =  features.shape[0] * features.shape[1]
                  cont = cont.add(Li)   # 1 value
              i+=1
</code></pre>
<p>I tried autograd backward but I get every time a RuntimeError. Then I saw that the grad_fn of Zi everytime remains None!!</p>

albanD: <p>Hi,</p>
<p>It is hard to say what the backward should be without knowing what the forward is.<br>
Also could you be more precise on what kind of error you’re seeing when trying to use your custom Function?</p>

Hdk: <p>Hi <a class="mention" href="/u/alband">@albanD</a> , thank you for reaching out.</p>
<p>I want to get these:<br>
<code>grad_weight += cont_loss_weight</code> and <code>grad_bias += cont_loss_bias</code>.</p>
<p>For this I have two ideas in my mind, however i don’t know which one is correct (If you can suggest!).<br>
Idea-1:</p>
<pre><code class="lang-auto"># keeping remaining code unchanged
if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output)  #shape = ([128,96,5,5])
            grad_feat = torch.autograd.grad(cont_loss, feat)         # has to be of shape = ([[batch_size= 100,128, 4, 4])
            cont_loss_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_feat )                                                                     
            grad_weight += cont_loss_weight

# But not sure how to get `cont_loss_bias` in same manner.

</code></pre>
<p>Idea-2:</p>
<pre><code class="lang-auto"># keeping remaining code unchanged
if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output)  #shape = ([128,96,5,5])
            cont_loss_weight = torch.autograd.grad(outputs= cont_loss, inputs= weight, retain_graph=(True))
            grad_weight += cont_loss_weight
            
if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum((0,2,3))        #shape = ([128])
            cont_loss_bias = torch.autograd.grad(outputs= cont_loss, inputs= bias, retain_graph=(True))
            grad_bias += cont_loss_bias          
         
</code></pre>
<p>Here’s the code :</p>
<pre><code class="lang-auto">class Custom_Convolution(torch.autograd.Function):    
    
    @staticmethod
    def forward(ctx, input, weight, bias, stride, padding):  #input(from previous layer)'s shape = ([batch_size=100, 96, 8, 8])
        output = torch.nn.functional.conv2d(input, weight, bias, stride, padding)  
        ctx.save_for_backward(input, weight, bias, output)
        return output    #output's shape = ([[batch_size= 100,128, 4, 4])

    @staticmethod
    def backward(ctx, grad_output):    # grad_output size = ([batch_size, 128,4,4])
        
        input, weight, bias, output = ctx.saved_tensors    #input size = ([batch_size, 96,8,8])
      
        feat = output.clone()#.requires_grad_(True)   # output from forward with size = ([batch_size, 128,4,4])

        feat = feat.view(feat.shape[0], feat.shape[1], -1) # features size = ([batch_size, 128,16])

        i = 0        
        cont = torch.tensor([0.]).to(dev)
        while i in range(0, feat.shape[0]):
              for f in range(len(feat[i])):
                  Zi_unnormalized = feat[i][f]
                  Zi = torch.nn.functional.normalize(Zi_unnormalized, dim = 0)
                  # Zj and Zk are tensors made from feat[i][*] and feat[other than i][*]. Zj and Zk varies for each Zi (or f)

                  Zi_Zk = torch.Tensor([0]).to(dev)
                  for k in Zk:
                      k= torch.nn.functional.normalize(k, dim = 0)
                      zi_zk = ...
                      Zi_Zk = Zi_Zk.add(zi_zk)

                  # Similarly computing Zi_Zj
                  # Li = some algebra of Zi_Zj and Zi_Zk
                  # number of 'Li' values =  feat.shape[0] * feat.shape[1]
                  cont = cont.add(Li)   # 1 value
              i+=1

        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output) #shape = ([batch_size,96,8,8])
        
        # If I go with Idea-2        
        if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output)  #shape = ([128,96,5,5])
            cont_loss_weight = torch.autograd.grad(outputs= cont_loss, 
                                            inputs= weight, retain_graph=(True))
            grad_weight += cont_loss_weight
            
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum((0,2,3))        #shape = ([128])
            cont_loss_bias = torch.autograd.grad(outputs= cont_loss, inputs= bias, retain_graph=(True))
           grad_bias += cont_loss_bias          
            
        if bias is not None:
            return grad_input, grad_weight, grad_bias, None, None
        else:
            return grad_input, grad_weight, None, None
 
</code></pre>
<p>It gives:<br>
<code>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</code></p>
<p>If I put <code>feat = output.clone().requires_grad_(True)</code>, it gives :</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
</code></pre>

albanD: <p>Hi,</p>
<p>Thanks for the code sample but you still didn’t say what you’re trying to accomplish here.</p>
<blockquote>
<p>RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</p>
</blockquote>
<p>This means that the output of your function does not require gradients. You need to make sure that at least one of the input Tensors requires gradients.</p>
<blockquote>
<p><code>feat = output.clone().requires_grad_(True)</code></p>
</blockquote>
<p>This would just make the output require gradients, that won’t make the autograd work with operations that happened before.<br>
You should have your input requiring gradients so that you can compute gradients for it.</p>

Hdk: <p>Let me break down the problem.</p>
<p>I have a network of say 5 convolutional layers, wherein only 4th layer is a custom convolution layer as you saw in above code.</p>
<p>The <code>output</code> shape from forward of this 4th i.e. custom layer is ([[batch_size,128,4,4]).<br>
Now, I have to include an additional loss (apart from cross entropy at the end of network which is usual in image classification task) in the backward computation of the layer itself, which takes in features (Zi, Zj, Zk) from that <code>output</code> to let the network learn them.</p>
<p>Sorry, there is one typo. Correction : <code>cont_loss = torch.tensor([0.]).to(dev)</code> and <code>cont_loss = cont_loss.add(Li)</code></p>
<p>cont_loss is that additional loss that you have noticed in the code. For including this only I have to manipulate the backward so that everything works as usual when I call<code> loss.backward()</code> in training loop.</p>
<p>That’s why I have these two additional terms which need to be taken care :<br>
<code>grad_weight += cont_loss_weight</code> and <code>grad_bias += cont_loss_bias</code>.</p>
<p>But, I don’t know which of my idea is correct that I wrote before.</p>
<p>I hope I could elaborate.</p>

albanD: <aside class="quote no-group quote-modified" data-username="Hdk" data-post="5" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/ec9cab/48.png" class="avatar"> Hdk:</div>
<blockquote>
<p>Now, I have to include an additional loss in the backward computation of the layer itself</p>
</blockquote>
</aside>
<p>Could you explain why you have to include this in the backward of the layer itself?<br>
Why can’t it be an additional term of your loss that you backprop as usual <code>(ce_loss + additional_loss).backward()</code> ?</p>

Hdk: <p>Well, that loss is specific to that layer only. I don’t have concrete argument for that but I was suggested to implement the loss within layer backward only. So, I have to do it that way (No option !<img src="https://discuss.pytorch.org/images/emoji/apple/zipper_mouth_face.png?v=9" title=":zipper_mouth_face:" class="emoji" alt=":zipper_mouth_face:">)</p>

albanD: <p>But if that loss term is only influenced by that layers weight, the gradient corresponding to that part of the loss will only influence that layer’s weights.</p>
<p>Also you still haven’t shared any formula of what you’re trying to compute so it is hard to say <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

Hdk: <p>Yes, you are right. It will only influence that layer’s weights.</p>
<p>Sorry, you can see equation of loss in equation no. (2) &amp; (3) and on page no. 13 the gradient computation.<br>
<a href="https://arxiv.org/pdf/2004.11362.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/2004.11362.pdf</a></p>

albanD: <p>For such problem, it will definitely be much easier to let the autograd figure out the gradients. In particular because you might not have all the information you need during the backward of that layer as the loss will depend on the layers that appeared after this one.</p>
<p>Note as well that a good ressource is other implementations of contrastive losses for pytorch: <a href="https://github.com/topics/contrastive-loss">https://github.com/topics/contrastive-loss</a></p>

Hdk: <p>As far as I saw other implementations, I found that a seperate loss function is created as you mentioned. But those implementations have two seperate networks one for contrastive loss and then fine tuning with the other one for classification task by CE loss.</p>
<p>But in my case I need to proceed with only one network with contrastive loss included in backward of one layer so that the network learns the features of that particular layer and should update its weight accordingly. That’s why there are these terms in backward:</p>
<p><code>grad_weight += cont_loss_weight</code> and <code>grad_bias += cont_loss_bias</code></p>
<p>So, I have to stick to this approach <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=9" title=":hugs:" class="emoji" alt=":hugs:"> but I haven’t seen such implementation or thread regarding this. <img src="https://discuss.pytorch.org/images/emoji/apple/man_shrugging.png?v=9" title=":man_shrugging:" class="emoji" alt=":man_shrugging:"></p>

Hdk: <aside class="quote no-group" data-username="albanD" data-post="6" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>(ce_loss + additional_loss).backward()</p>
</blockquote>
</aside>
<p>In doing so, <a href="https://discuss.pytorch.org/t/custom-loss-functions/29387/60">this</a> will happen as mentioned by <a class="mention" href="/u/ptrblck">@ptrblck</a>. So, the layers before the custom layers will also have gradients of contrastive loss accumulated in addition to that of CE loss. But, I want to restrict this gradient accumulation (of contrastive + CE loss) within the custom layer.</p>
<aside class="quote no-group" data-username="albanD" data-post="6" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Why can’t it be an additional term of your loss that you backprop as usual</p>
</blockquote>
</aside>
<p>Probably, this would justify why I am doing it this way. <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>

albanD: <p>In that case, you can use the nightly version of pytorch an use the new <code>inputs</code> argument to the <a href="https://pytorch.org/docs/master/autograd.html?highlight=backward#torch.autograd.backward">.backward()</a> function:</p>
<pre><code class="lang-python">net.zero_grad()
ce_loss.backward()
additional_loss.backward(inputs=net.your_contrastive_layer.parameters())
opt.step()
</code></pre>

Hdk: <aside class="quote no-group" data-username="albanD" data-post="13" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>In that case, you can use the nightly version of pytorch an use the new <code>inputs</code> argument to the <a href="https://pytorch.org/docs/master/autograd.html?highlight=backward#torch.autograd.backward" rel="noopener nofollow ugc">.backward()</a> function:</p>
</blockquote>
</aside>
<p>So for doing this you mean, I should create a seperate Contrastive loss function instead of doing it in layer backward. And then use nightly version code that you mentioned. Right?</p>
<p>And can you please explain what this code will do. <code>additional_loss.backward(inputs=net.your_contrastive_layer.parameters()) </code></p>
<p>I try to install nightly version with from <a href="https://anaconda.org/pytorch/pytorch-nightly" rel="noopener nofollow ugc">here</a>.<br>
But I encounterd an error this error:</p>
<pre><code class="lang-auto">EnvironmentNotWritableError: The current user does not have write permissions to the target environment.
  environment location: C:\ProgramData\Anaconda3
</code></pre>
<p>Looks like I don’t have permission from the admin of the PC I’m using.</p>

albanD: <blockquote>
<p>So for doing this you mean, I should create a seperate Contrastive loss function instead of doing it in layer backward.</p>
</blockquote>
<p>I think it is going to be simpler than modifying the backward and writing the gradients yourself yes.</p>
<blockquote>
<p>And can you please explain what this code will do.</p>
</blockquote>
<p>It will run the backward as usual but will only update the <code>.grad</code> fields of the <code>inputs</code> that were given. So in your case, you only want to update the parameters of that one layer.</p>
<blockquote>
<p>Looks like I don’t have permission from the admin of the PC I’m using.</p>
</blockquote>
<p>You might want to create a new environment in your conda so that you can install things.</p>

Hdk: <p>Thanks for the explanation. <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="albanD" data-post="15" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>You might want to create a new environment in your conda so that you can install things.</p>
</blockquote>
</aside>
<p>I will have to look into this.</p>
<p>But, <strong>if</strong> I stick to <a href="https://discuss.pytorch.org/t/gradient-computation-in-custom-backward/102810/3">this</a> code above. What do you think is the correct way to calculate <code>cont_loss_weight</code> by Idea-1 or Idea-2?</p>
<p>And what changes should I make in the code of backward to run it successfully, any suggestions?</p>
<p>I tried this as well without nightly version, but GPU went out of memory. <img src="https://discuss.pytorch.org/images/emoji/apple/persevere.png?v=12" title=":persevere:" class="emoji" alt=":persevere:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="albanD" data-post="6" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>(ce_loss + additional_loss).backward()</p>
</blockquote>
</aside>
<p>RuntimeError: CUDA out of memory.</p>

albanD: <blockquote>
<p>But, <strong>if</strong> I stick to <a href="https://discuss.pytorch.org/t/gradient-computation-in-custom-backward/102810/3">this</a> code above. What do you think is the correct way to calculate <code>cont_loss_weight</code> by Idea-1 or Idea-2?</p>
</blockquote>
<p>I honestly don’t know. You will need to derive the mathematical formula for what the gradient should be with pen and paper first. Then implement the final formula you get in there.</p>
<blockquote>
<p>I tried this as well without nightly version, but GPU went out of memory. <img src="https://discuss.pytorch.org/images/emoji/apple/persevere.png?v=12" title=":persevere:" class="emoji" alt=":persevere:" loading="lazy" width="20" height="20"></p>
</blockquote>
<p>Does it run out at the first iteration? Or after a while?<br>
Can you try reducing the batch size to reduce memory pressure?</p>

Hdk: <aside class="quote no-group" data-username="albanD" data-post="17" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>You will need to derive the mathematical formula for what the gradient should be with pen and paper first.</p>
</blockquote>
</aside>
<p>In the <a href="https://arxiv.org/pdf/2004.11362.pdf" rel="noopener nofollow ugc">research paper</a> that I showed you, there is gradient equation (i.e. differentiation of Li w.r.t. Zi for every 'i’th feature if you compare with my code), but I don’t know why it was w.r.t. evey feature instead of w.r.t. parameters (weight and bias). Contradictorily, the loss formula doesn’t include use of parameters, but the features which go into the loss formula are result of mathematics of input and parameters!</p>
<aside class="quote no-group" data-username="albanD" data-post="17" data-topic="102810">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Does it run out at the first iteration? Or after a while?<br>
Can you try reducing the batch size to reduce memory pressure?</p>
</blockquote>
</aside>
<p>If I use smaller batch size then I could see results for first few batches, but I doubt that it would still go out of memory after some more iterations.<br>
However, this loss function yields good results for bigger batch sizes. So, I tried with bigger batch size,but it went out of memory after some time without showing first result.</p>

Hdk: <p>I’m witnessing some unusual output for below code, can you have a look please!</p>
<pre><code class="lang-auto">class Custom_Convolution(torch.autograd.Function):    
    
    @staticmethod
    def forward(ctx, input, weight, bias, stride, padding):  #input(from previous layer)'s shape = ([batch_size=100, 96, 8, 8])
        output = torch.nn.functional.conv2d(input, weight, bias, stride, padding)  
        ctx.save_for_backward(input, weight, bias, output)
        return output    #output's shape = ([[batch_size= 100,128, 4, 4])

    @staticmethod
    def backward(ctx, grad_output):    # grad_output size = ([batch_size, 128,4,4])
        
        input, weight, bias, output = ctx.saved_tensors    #input size = ([batch_size, 96,8,8])  
        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output) #shape = ([batch_size,96,8,8])
              
        if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output)  #shape = ([128,96,5,5])
                        
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum((0,2,3))        #shape = ([128])          
           
        with torch.enable_grad():
               feat = output.clone()   # output from forward with size = ([batch_size, 128,4,4])

               feat = feat.view(feat.shape[0], feat.shape[1], -1) # features size = ([batch_size, 128,16])
       
               cont = torch.tensor([0.]).to(dev)
               for i in range(0, feat.shape[0]):
                   for f in range(len(feat[i])):
                       Zi_unnormalized = feat[i][f]
                       Zi = torch.nn.functional.normalize(Zi_unnormalized, dim = 0)
                       # Zj and Zk are tensors made from feat[i][*] and feat[other than i][*]. Zj and Zk varies for each Zi (or f)

                       Zi_Zk = torch.Tensor([0]).to(dev)
                       for k in Zk:
                           k= torch.nn.functional.normalize(k, dim = 0)
                           zi_zk = ...
                           Zi_Zk = Zi_Zk.add(zi_zk)

                       # Similarly computing Zi_Zj
                       # Li = some algebra of Zi_Zj and Zi_Zk
                       # number of 'Li' values =  feat.shape[0] * feat.shape[1]
                       cont = cont.add(Li)   # 1 value
               print("\n Loss: ", cont_loss, cont_loss.requires_grad)

### This line of printing loss keeps on repeating with the same value of Loss!!!! 


        cont_loss_weight = torch.autograd.grad(outputs= cont_loss, inputs= weight, retain_graph=(True))
        print ("Shape:", cont_loss_weight .shape)                                   
        grad_weight += cont_loss_weight

        cont_loss_bias = torch.autograd.grad(outputs= cont_loss, inputs= bias, retain_graph=(True))
        grad_bias += cont_loss_bias

        if bias is not None:
            return grad_input, grad_weight, grad_bias, None, None
        else:
            return grad_input, grad_weight, None, None

</code></pre>
<p>Output :</p>
<pre><code class="lang-auto">Loss:  tensor([37.218], device='cuda:0', grad_fn=&lt;AddBackward0&gt;) True
Loss:  tensor([37.218], device='cuda:0', grad_fn=&lt;AddBackward0&gt;) True
Loss:  tensor([37.218], device='cuda:0', grad_fn=&lt;AddBackward0&gt;) True
Loss:  tensor([37.218], device='cuda:0', grad_fn=&lt;AddBackward0&gt;) True
.
.
.
RuntimeError: CUDA out of memory.
</code></pre>
<p>I don’t know why this line is repeating so many time and has no end. And finally CUDA goes out of memory. It should print once only for 1 batch, moreover I haven’t made any indentation mistake!</p>
<p>This line<br>
<code>cont_loss_weight = torch.autograd.grad(outputs= cont_loss, inputs= weight, retain_graph=(True))</code><br>
is getting executed but it neither show any error nor it returns something because its following line: <code>print ("Shape:", cont_loss_weight .shape)</code> doesn’t get printed.</p>

Hdk: <p>Hi <a class="mention" href="/u/alband">@albanD</a>,</p>
<p>When I tried to include autograd.grad in backward as <a href="https://discuss.pytorch.org/t/gradient-computation-in-custom-backward/102810/19">above</a>, autograd.grad wasn’t returning anything though is was getting executed. I don’t know why, can you please have a look!</p>
<p>Then I tried it with different approach :</p>
<pre><code class="lang-auto">class Custom_Convolution(torch.autograd.Function):    
    
    @staticmethod
    def forward(ctx, input, weight, bias, stride, padding):  #input(from previous layer)'s shape = ([batch_size=100, 96, 8, 8])
        with torch.enable_grad():
               output = torch.nn.functional.conv2d(input, weight, bias, stride, padding)
               h = output.shape[2]
               w = output.shape[3]  
               # output from forward with size = ([batch_size, 128,4,4])

               output= output.view(output.shape[0], output.shape[1], -1) # output size = ([batch_size, 128,16])
       
               cont = torch.tensor([0.]).to(dev).requires_grad_(True)
               for i in range(0, output.shape[0]):
                   for f in range(len(output[i])):
                       Zi_unnormalized = output[i][f]
                       Zi = torch.nn.functional.normalize(Zi_unnormalized, dim = 0)
                       # Zj and Zk are tensors made from output[i][*] and output[other than i][*]. Zj and Zk varies for each Zi (or f)

                       Zi_Zk = torch.Tensor([0]).to(dev).requires_grad_(True)
                       for k in Zk:
                           k= torch.nn.functional.normalize(k, dim = 0)
                           zi_zk = ...
                           Zi_Zk = Zi_Zk.add(zi_zk)

                       # Similarly computing Zi_Zj
                       # Li = some algebra of Zi_Zj and Zi_Zk
                       # number of 'Li' values =  output.shape[0] * output.shape[1]
                       cont = cont.add(Li)   # 1 value
        print("\n Loss: ", cont_loss, cont_loss.requires_grad)
           
        # weight1 = weight.clone().requires_grad_(True)
        # bias1 = bias.clone().requires_grad_(True)

        # weight.shape = ([128, 96, 5, 5])
        cont_loss_weight  = torch.autograd.grad(outputs= cont_loss,inputs= weight, retain_graph=True)
    
        #bias.shape = ([128])
        cont_loss_bias = torch.autograd.grad(outputs= cont_loss, inputs= bias, retain_graph=True)
        
        output = output.view(output.shape[0], output.shape[1], h,w)
        ctx.save_for_backward(input, weight, bias, output, cont_loss, cont_loss_weight,cont_loss_bias)            

        return output    #output's shape = ([[batch_size= 100,128, 4, 4])

    @staticmethod
    def backward(ctx, grad_output):    # grad_output size = ([batch_size, 128,4,4])
        
        input, weight, bias, output,  cont_loss,cont_loss_weight,cont_loss_bias = ctx.saved_tensors    #input size = ([batch_size, 96,8,8])  
        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output) #shape = ([batch_size,96,8,8])
              
        if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output)  #shape = ([128,96,5,5])
            grad_weight += cont_loss_weight 
                        
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum((0,2,3))        #shape = ([128])         
            grad_bias += cont_loss_bias

        if bias is not None:
            return grad_input, grad_weight, grad_bias, None, None
        else:
            return grad_input, grad_weight, None, None
</code></pre>
<p>Then I observed that <code>cont_loss_weight</code> is a tuple object containing two tensors each of shape ([96, 5, 5]). It should have returned a tensor of shape ([128, 96, 5, 5]) instead of tuple. And similarly for <code>cont_loss_bias</code>, a tensor of shape  ([128]).<br>
I don’t know why!</p>
<p>Moreover, when I do `cont_loss_weight  = torch.autograd.grad(outputs= cont_loss,inputs= weight, retain_graph=True), I am guessing grad_weight in backward will get affected. I have to keep retain_graph= True as well.</p>
<p>So to avoid that, when I used a copy of parameters i.e. <code>cont_loss_weight  = torch.autograd.grad(outputs= cont_loss,inputs= weight1, retain_graph=True)</code> I got this error</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
</code></pre>