Solved by albanD in post #2 
              
              
                Hi, 
You don’t need to use Variable anymore. You can call .requires_grad_() on your Tensor after building it. 
The problem here I think is your “SaveFeatures”. You rewrap the output in a brand new Tensor. That means that you break the link with the existing graph. 
Why not do self.features = output …