Solved by ptrblck in post #2 
              
              
                By rewrapping your intermediate activation in a new tensor, you are detaching it from the computation graph: 
self.features = torch.tensor(output, requires_grad=True, device=device)

You could assign self.features directly to output, which will give you a valid gradient in your input tensors.