Solved by albanD in post #3 
              
              
                Hi, 
In pytorch, you compute what you need to get the output in the forward pass and it will just backward on that part. 
So yes you code will do what you want.